cleaning

src/anemoi/training/diagnostics/callbacks/checkpoint.py

@@ -77,20 +77,11 @@ def model_metadata(self, model: torch.nn.Module) -> dict:
 
         return self._model_metadata
 
-    def on_validation_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
-        """Save a checkpoint at the end of the validation stage."""
-        del pl_module
-        if not self._should_skip_saving_checkpoint(trainer) and not self._should_save_on_train_epoch_end(trainer):
-            monitor_candidates = self._monitor_candidates(trainer)
-            if self._every_n_epochs >= 1 and (trainer.current_epoch + 1) % self._every_n_epochs == 0:
-                self._save_topk_checkpoint(trainer, monitor_candidates)
-            self._save_last_checkpoint(trainer, monitor_candidates)
-
     def on_fit_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
         del pl_module
         if not self._should_skip_saving_checkpoint(trainer) and not self._should_save_on_train_epoch_end(trainer):
             monitor_candidates = self._monitor_candidates(trainer)
-            # Need to correct the checkpoint epoch to the last epoch
+            # PTL advances one epoch at end of training,  Need to correct the checkpoint epoch to the last epoch
             monitor_candidates["epoch"] = trainer.current_epoch - 1
             self._save_topk_checkpoint(trainer, monitor_candidates)
             self._save_last_checkpoint(trainer, monitor_candidates)