Fix UnboundLocalError when using the predict method with return_predi…

…ctions=False. (#20484) * any_on_epoch referenced before assignment * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Luca Antiga <[email protected]>
Lightning-AI · Dec 11, 2024 · 1e32ebf · 1e32ebf
1 parent 601c060
commit 1e32ebf
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/lightning/pytorch/loops/prediction_loop.py b/src/lightning/pytorch/loops/prediction_loop.py
@@ -233,8 +233,9 @@ def _predict_step(
 
         self.batch_progress.increment_ready()
 
-        if not using_dataloader_iter:
-            any_on_epoch = self._store_data_for_prediction_writer(batch_idx, dataloader_idx)
+        any_on_epoch = (
+            self._store_data_for_prediction_writer(batch_idx, dataloader_idx) if not using_dataloader_iter else False
+        )
 
         # the `_step` methods don't take a batch_idx when `dataloader_iter` is used, but all other hooks still do,
         # so we need different kwargs