Merge pull request #23 from moskomule/dev

fix weight save
moskomule · Jul 2, 2019 · be11001 · be11001
2 parents 7e78f45 + 22af20c
commit be11001
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/homura/callbacks/saver.py b/homura/callbacks/saver.py
@@ -40,10 +40,16 @@ def save(self,
              data: Mapping,
              file_name: str):
         try:
+            # scheduler is not a must
+            scheduler_state_dict = data.get(SCHEDULER)
+            if scheduler_state_dict is not None:
+                scheduler_state_dict = scheduler_state_dict.state_dict()
+
             torch.save({"git": get_git_hash(),
                         "args": get_args(),
                         MODEL: data[MODEL].state_dict(),
                         OPTIMIZER: data[OPTIMIZER].state_dict(),
+                        SCHEDULER: scheduler_state_dict,
                         EPOCH: self._epoch,
                         STEP: self._step},
                        self.save_path / file_name)