Fix AttributeError when using CombinedLoader in prediction

CHANGELOG.md

@@ -290,10 +290,12 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Avoid the deprecated `onnx.export(example_outputs=...)` in torch 1.10 ([#11116](https://github.com/PyTorchLightning/pytorch-lightning/pull/11116))
 
 
-
 - Fixed an issue when torch-scripting a `LightningModule` after training with `Trainer(sync_batchnorm=True)` ([#11078](https://github.com/PyTorchLightning/pytorch-lightning/pull/11078))
 
 
+- Fixed an `AttributeError` occuring when using a `CombinedLoader` (multiple dataloaders) for prediction ([#11111](https://github.com/PyTorchLightning/pytorch-lightning/pull/11111))
+
+
 - Fixed bug where `Trainer(track_grad_norm=..., logger=False)' would fail ([#11114](https://github.com/PyTorchLightning/pytorch-lightning/pull/11114))
 
 

pytorch_lightning/loops/epoch/prediction_epoch_loop.py

@@ -162,7 +162,8 @@ def _build_kwargs(self, batch: Any, batch_idx: int, dataloader_idx: int) -> Dict
     def _get_batch_indices(self, dataloader_idx: int) -> List[List[int]]:
         """Returns a reference to the seen batch indices if the dataloader has a batch sampler wrapped by our
         :class:`~pytorch_lightning.overrides.distributed.IndexBatchSamplerWrapper`."""
-        batch_sampler = self.trainer.predict_dataloaders[dataloader_idx].batch_sampler
+        # the batch_sampler is not be defined in case of CombinedDataLoaders
+        batch_sampler = getattr(self.trainer.predict_dataloaders[dataloader_idx], "batch_sampler", None)
         if isinstance(batch_sampler, IndexBatchSamplerWrapper) and self.should_store_predictions:
             return batch_sampler.seen_batch_indices
 

tests/callbacks/test_prediction_writer.py

@@ -16,8 +16,10 @@
 import pytest
 from torch.utils.data import DataLoader
 
+import pytorch_lightning as pl
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import BasePredictionWriter
+from pytorch_lightning.trainer.supporters import CombinedLoader
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.helpers import BoringModel, RandomDataset
 from tests.helpers.runif import RunIf
@@ -106,3 +108,35 @@ def test_prediction_writer_batch_indices(tmpdir, num_workers):
             call(trainer, model, ANY, [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]]),
         ]
     )
+
+
+def test_prediction_writer_partial_support_for_combined_loader(tmpdir):
+    """Test partial support for CombinedLoader: prediction works but sample indices don't get tracked."""
+    pl.loops.epoch.prediction_epoch_loop.warning_cache.clear()
+
+    class PredictionModel(BoringModel):
+        def predict_dataloader(self):
+            return CombinedLoader(
+                {
+                    "a": DataLoader(RandomDataset(32, 8), batch_size=2),
+                    "b": DataLoader(RandomDataset(32, 8), batch_size=4),
+                }
+            )
+
+        def predict_step(self, batch, *args, **kwargs):
+            return self(batch["a"])
+
+    DummyPredictionWriter.write_on_batch_end = Mock()
+    DummyPredictionWriter.write_on_epoch_end = Mock()
+
+    model = PredictionModel()
+    writer = DummyPredictionWriter("batch_and_epoch")
+    trainer = Trainer(callbacks=writer)
+    with pytest.warns(UserWarning, match="Lightning couldn't infer the indices fetched for your dataloader."):
+        trainer.predict(model)
+
+    writer.write_on_batch_end.assert_has_calls(
+        [call(trainer, model, ANY, [], ANY, 0, 0), call(trainer, model, ANY, [], ANY, 1, 0)]
+    )
+
+    writer.write_on_epoch_end.assert_has_calls([call(trainer, model, ANY, [[]])])