Merge pull request #160 from delira-dev/save_fix

Fix Epoch during saving
delira-dev · Jul 13, 2019 · db58949 · db58949
2 parents 722a536 + 3b7388b
commit db58949
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/delira/training/pytorch_trainer.py b/delira/training/pytorch_trainer.py
@@ -436,7 +436,7 @@ def save_state(self, file_name, epoch, **kwargs):
             """
             if not (file_name.endswith(".pth") or file_name.endswith(".pt")):
                 file_name = file_name + ".pt"
-            save_checkpoint(file_name, self.module, self.optimizers,
+            save_checkpoint(file_name, self.module, self.optimizers, epoch,
                             **kwargs)
 
         @staticmethod