Update main.py

agosztolai · Dec 6, 2023 · 5fa052e · 5fa052e
1 parent cb75f38
commit 5fa052e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/MARBLE/main.py b/MARBLE/main.py
@@ -413,7 +413,8 @@ def load_model(self, loadpath):
         self._epoch = checkpoint["epoch"]
         self.load_state_dict(checkpoint["model_state_dict"])
         self.optimizer_state_dict = checkpoint["optimizer_state_dict"]
-        self.losses = checkpoint["losses"]
+        if hatattr(self, 'losses'):
+            self.losses = checkpoint["losses"]
 
     def save_model(self, optimizer, losses, outdir=None, best=False, timestamp=""):
         """Save model."""