fix a minor bug with async checkpointing where a checkpoint would get…

… saved on_train_batch_end and on_validation_end within the same step (NVIDIA#9856) (NVIDIA#9867)

Signed-off-by: ashors1 <[email protected]>
Co-authored-by: Anna Shors <[email protected]>

nemo/lightning/pytorch/callbacks/model_checkpoint.py

@@ -380,6 +380,8 @@ def _save_checkpoint(self, trainer: 'pytorch_lightning.Trainer', filepath: str)
         self.set_checkpoint_unfinished_marker(filepath, barrier_after=True)
         ema_callback = self._ema_callback(trainer)
 
+        self._last_global_step_saved = trainer.global_step
+
         if ema_callback is not None:
             if self.async_save:
                 raise ValueError('async_save with EMA not supported')
@@ -422,7 +424,6 @@ def _get_finalize_save_checkpoint_callback(
 
         def _cb():
             logging.debug(f'Finalize callback called for step {global_step}, filepath {filepath}')
-            self._last_global_step_saved = global_step
             self._last_checkpoint_saved = filepath
 
             from nemo.utils.get_rank import is_global_rank_zero