Suppress Warning in PredictionEpochLoop (#11189)

Co-authored-by: Rohit Gupta <[email protected]> Co-authored-by: Carlos Mocholí <[email protected]>
Lightning-AI · Dec 21, 2021 · 98fcf1e · 98fcf1e
1 parent 1f32923
commit 98fcf1e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pytorch_lightning/loops/epoch/prediction_epoch_loop.py b/pytorch_lightning/loops/epoch/prediction_epoch_loop.py
@@ -68,8 +68,9 @@ def on_run_start(
         void(dataloader_iter, dataloader_idx)
         self._dl_max_batches = dl_max_batches
         self._num_dataloaders = num_dataloaders
-        self._seen_batch_indices = self._get_batch_indices(dataloader_idx)
         self.return_predictions = return_predictions
+        # this call requires that `self.return_predictions` is set
+        self._seen_batch_indices = self._get_batch_indices(dataloader_idx)
 
     def advance(
         self,