reset progress for restart

Lightning-AI · Aug 10, 2021 · f69dbc4 · f69dbc4
1 parent 3096ab8
commit f69dbc4
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/pytorch_lightning/loops/epoch/training_epoch_loop.py b/pytorch_lightning/loops/epoch/training_epoch_loop.py
@@ -229,6 +229,10 @@ def on_run_end(self) -> List[List[STEP_OUTPUT]]:
 
         self.update_lr_schedulers("epoch", update_plateau_schedulers=True)
 
+        self.batch_progress.current.reset()
+        self.scheduler_progress.current.reset()
+        self.batch_loop.optim_progress.reset_on_epoch()
+
         epoch_output = self._epoch_output
         # free memory
         self._epoch_output = None