unwrap model and optimizer manually during save and load for fabric

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -24,6 +24,7 @@
 from ema_pytorch import EMA
 
 from lightning import Fabric
+from lightning.fabric.wrappers import _unwrap_objects
 
 # constants
 
@@ -341,9 +342,12 @@ def save(
 
         path.parent.mkdir(exist_ok = True, parents = True)
 
+        unwrapped_model = _unwrap_objects(self.model)
+        unwrapped_optimizer = _unwrap_objects(self.optimizer)
+
         package = dict(
-            model = self.model.state_dict_with_init_args,
-            optimizer = self.optimizer.state_dict(),
+            model = unwrapped_model.state_dict_with_init_args,
+            optimizer = unwrapped_optimizer.state_dict(),
             scheduler = self.scheduler.state_dict(),
             steps = self.steps
         )
@@ -379,9 +383,14 @@ def load(
 
         self.model_loaded_from_path = path
 
+        # get unwrapped model and optimizer
+
+        unwrapped_model = _unwrap_objects(self.model)
+        unwrapped_optimizer = _unwrap_objects(self.optimizer)
+
         # load model from path
 
-        self.model.load(path)
+        unwrapped_model.load(path)
 
         if only_model:
             return
@@ -391,7 +400,7 @@ def load(
         package = torch.load(str(path))
 
         if 'optimizer' in package:
-            self.optimizer.load_state_dict(package['optimizer'])
+            unwrapped_optimizer.load_state_dict(package['optimizer'])
 
         if 'scheduler' in package:
             self.scheduler.load_state_dict(package['scheduler'])

pyproject.toml

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