Update checkpointer.py

vturrisi · Jun 18, 2024 · ccdba0b · ccdba0b
1 parent 3d58244
commit ccdba0b
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/solo/utils/checkpointer.py b/solo/utils/checkpointer.py
@@ -174,7 +174,12 @@ def on_train_epoch_end(self, trainer: pl.Trainer, _):
         if epoch % self.frequency == 0:
             self.save(trainer)
 
-    def on_train_epoch_end(self, *args, **kwargs):
-        """Saves model at the end of training. """
+    def on_train_end(self, trainer: pl.
+Trainer, _):
+        """Saves model at the end of training.
+
+        Args:
+            trainer (pl.Trainer): pytorch lightning trainer object.
+        """
 
         self.save(trainer)