Always save model at the end of training

solo/utils/checkpointer.py

@@ -164,7 +164,7 @@ def on_train_start(self, trainer: pl.Trainer, _):
         self.save_args(trainer)
 
     def on_train_epoch_end(self, trainer: pl.Trainer, _):
-        """Tries to save current checkpoint at the end of each train epoch.
+        """Tries to save the current checkpoint at the end of each train epoch.
 
         Args:
             trainer (pl.Trainer): pytorch lightning trainer object.
@@ -173,3 +173,8 @@ def on_train_epoch_end(self, trainer: pl.Trainer, _):
         epoch = trainer.current_epoch  # type: ignore
         if epoch % self.frequency == 0:
             self.save(trainer)
+
+    def on_train_epoch_end(self, *args, **kwargs):
+        """Saves model at the end of training. """
+
+        self.save(trainer)