code: def test_prediction_loop_batch_sampler_set_epoch_called(
      tmp_path, use_distributed_sampler
  ):
      """Tests that set_epoch is called on the dataloader's batch sampler (if any) during prediction."""
      trainer = Trainer(
          default_root_dir=tmp_path,
          limit_predict_batches=1,
          enable_model_summary=False,
          enable_checkpointing=False,
          logger=False,
          strategy="ddp",
          devices=1,
          accelerator="cpu",
          use_distributed_sampler=use_distributed_sampler,
      )
  
      class MyModel(BoringModel):
          def predict_dataloader(self):
              dataset = RandomDataset(32, 64)
              sampler = None
              if not use_distributed_sampler:
                  sampler = DistributedSampler(dataset)
              return DataLoader(dataset, sampler=sampler)
  
      model = MyModel()
      trainer.fit_loop.epoch_progress.current.processed = 2
      trainer.predict(model)
  
      # torch will set this .sampler attribute for backwards compatibility, but in reality, the batch sampler is used
      assert isinstance(
          trainer.predict_dataloaders.sampler, SequentialSampler
      )
      batch_sampler = trainer.predict_dataloaders.batch_sampler
      assert isinstance(batch_sampler, _IndexBatchSamplerWrapper)
      assert isinstance(batch_sampler.sampler, DistributedSampler)
      assert batch_sampler.sampler.epoch == 2
  
file_path: pytorch-lightning/tests/tests_pytorch/loops/test_prediction_loop.py
imports: import itertools
  import pytest
  from lightning.pytorch import LightningModule, Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
  from lightning.pytorch.overrides.distributed import _IndexBatchSamplerWrapper
  from torch.utils.data import DataLoader, DistributedSampler, SequentialSampler
label_desc: Use image node_id 2 for calling a global function with example usage: test_prediction_loop_batch_sampler_set_epoch_called(tmp_path, use_distributed_sampler) and returns: DataLoader