Fix bug in load_model in training

close LeeDoYup#12
stablegradients · Jan 30, 2021 · 0f22e7f · 0f22e7f
1 parent cff901f
commit 0f22e7f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/models/fixmatch/fixmatch.py b/models/fixmatch/fixmatch.py
@@ -266,6 +266,10 @@ def load_model(self, load_path):
                     eval_model.load_state_dict(checkpoint[key])
                 elif key == 'it':
                     self.it = checkpoint[key]
+                elif key == 'scheduler':
+                    self.scheduler.load_state_dict(checkpoint[key])
+                elif key == 'optimizer':
+                    self.optimizer.load_state_dict(checkpoint[key]) 
                 else:
                     getattr(self, key).load_state_dict(checkpoint[key])
                 self.print_fn(f"Check Point Loading: {key} is LOADED")