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
Tests that set_epoch is called on the dataloader's batch sampler (if any) during prediction.
tests/tests_pytorch/loops/test_prediction_loop.py