Handle end of batch

Author: lantiga

src/lightning/pytorch/loops/training_epoch_loop.py

@@ -150,6 +150,12 @@ def reset(self) -> None:
             # when loop state dict is saved on_train_batch_end, that is, before increment completed is called
             if self.batch_progress.total.completed < self.batch_progress.total.processed:
                 self.batch_progress.increment_completed()
+                # handle situation in which save happened on_train_batch_end and epoch is at end
+                if self.batch_progress.current.completed >= self.trainer.num_training_batches:
+                    self.batch_progress.reset_on_run()
+                    self.scheduler_progress.reset_on_run()
+                    self.automatic_optimization.optim_progress.reset_on_run()
+                    self.val_loop.batch_progress.total.reset()
                 if not self._should_accumulate():
                     self._batches_that_stepped += 1