Trainer train simplified, tests for load (#168)

- the arguments of Trainer.train now are passed to the fit - unittest for load/restoring from checkpoint
2023-07-25 17:23:12 +02:00
parent de0c3fca82
commit e84def3bf9
3 changed files with 44 additions and 7 deletions
--- a/tests/test_solvers/test_garom.py
+++ b/tests/test_solvers/test_garom.py
@@ -134,7 +134,7 @@ def test_train_cpu():
                                                hidden_dimension=64)
                )

-    trainer = Trainer(solver=solver, kwargs={'max_epochs' : 4, 'accelerator': 'cpu'})
+    trainer = Trainer(solver=solver, max_epochs=4, accelerator='cpu')
    trainer.train()

 def test_sample():