[loops] Reset reference to dataloader iterator on run end

Lightning-AI · Sep 9, 2021 · cd766b8 · cd766b8
1 parent 41ba639
commit cd766b8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py b/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py
@@ -135,6 +135,7 @@ def on_run_end(self) -> EPOCH_OUTPUT:
         outputs = self.outputs
         # free memory
         self.outputs = []
+        self.dataloader_iter = None
         return outputs
 
     def evaluation_step(self, batch: Any, batch_idx: int, dataloader_idx: int) -> Optional[STEP_OUTPUT]:

diff --git a/pytorch_lightning/loops/epoch/training_epoch_loop.py b/pytorch_lightning/loops/epoch/training_epoch_loop.py
@@ -232,6 +232,8 @@ def on_run_end(self) -> None:
         if self._num_training_batches_reached(self.is_last_batch):
             self.update_lr_schedulers("epoch", update_plateau_schedulers=True)
 
+        self.dataloader_iter = None
+
     def teardown(self) -> None:
         self._results.cpu()
         self.batch_loop.teardown()