fix trainer save / load

alphafold3_pytorch/trainer.py

@@ -215,14 +215,16 @@ def save(self, path: str | Path, overwrite = False):
             steps = self.steps
         )
 
-        torch.save(str(path), package)
+        torch.save(package, str(path))
 
     def load(self, path: str | Path, strict = True):
         if isinstance(path, str):
             path = Path(path)
 
         assert path.exists()
 
+        self.model.load(path)
+
         package = torch.load(str(path))
 
         if 'optimizer' in package:
@@ -233,8 +235,6 @@ def load(self, path: str | Path, strict = True):
 
         self.steps = package.get('steps', 0)
 
-        self.model.load_state_dict(package['model'])
-
     # shortcut methods
 
     def wait(self):

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.0.59"
+version = "0.0.60"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_trainer.py

@@ -139,7 +139,7 @@ def test_trainer():
 
     # saving and loading from trainer
 
-    trainer.save('./some/nested/folder2/training')
+    trainer.save('./some/nested/folder2/training', overwrite = True)
     trainer.load('./some/nested/folder2/training')
 
     # also allow for loading Alphafold3 directly from training ckpt