Make _batches_that_stepped consistent with progress

Author: lantiga

src/lightning/pytorch/loops/training_epoch_loop.py

@@ -270,6 +270,10 @@ def advance(self, data_fetcher: _DataFetcher) -> None:
         # failure to do this will lead to incorrect total batch completed progress counts
         self.batch_progress.increment_completed()
 
+        if not self._should_accumulate():
+            # this is increased once per batch disregarding multiple optimizers on purpose for loggers
+            self._batches_that_stepped += 1
+
         call._call_callback_hooks(trainer, "on_train_batch_end", batch_output, batch, batch_idx)
         call._call_lightning_module_hook(trainer, "on_train_batch_end", batch_output, batch, batch_idx)
         trainer._logger_connector.on_batch_end()
@@ -301,9 +305,6 @@ def on_advance_end(self, data_fetcher: _DataFetcher) -> None:
         # update plateau LR scheduler after metrics are logged
         self.update_lr_schedulers("step", update_plateau_schedulers=True)
 
-        if not self._should_accumulate():
-            # this is increased once per batch disregarding multiple optimizers on purpose for loggers
-            self._batches_that_stepped += 1
         # this will save based on the `batches_that_stepped` value
         self._save_loggers_on_train_batch_end()
 

tests/tests_pytorch/loops/test_loops.py

@@ -575,6 +575,8 @@ def test_fit_loop_reset(tmp_path):
     assert epoch_loop.batch_progress.current.processed == 2
     assert epoch_loop.batch_progress.current.completed == 2
 
+    assert epoch_loop._batches_that_stepped == 2
+
     assert optimizer_loop.restarting
 
     # reset state loaded from a checkpoint from the end of an epoch
@@ -605,6 +607,8 @@ def test_fit_loop_reset(tmp_path):
     assert epoch_loop.batch_progress.current.processed == 4
     assert epoch_loop.batch_progress.current.completed == 4
 
+    assert epoch_loop._batches_that_stepped == 4
+
 
 @pytest.mark.parametrize(
     ("train_datasets", "val_datasets"),