def test_evaluation_loop_sampler_set_epoch_called(tmp_path, use_batch_sampler):
    """Tests that set_epoch is called on the dataloader's sampler (if any) during training and validation."""

    def _get_dataloader():
        dataset = RandomDataset(32, 64)
        sampler = RandomSampler(dataset)
        sampler.set_epoch = Mock()
        if use_batch_sampler:
            batch_sampler = BatchSampler(sampler, 2, True)
            return DataLoader(dataset, batch_sampler=batch_sampler)
        return DataLoader(dataset, sampler=sampler)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=1,
        limit_val_batches=1,
        max_epochs=2,
        enable_model_summary=False,
        enable_checkpointing=False,
        logger=False,
    )

    train_dataloader = _get_dataloader()
    val_dataloader = _get_dataloader()
    trainer.fit(model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
    train_sampler = train_dataloader.batch_sampler.sampler if use_batch_sampler else train_dataloader.sampler
    val_sampler = val_dataloader.batch_sampler.sampler if use_batch_sampler else val_dataloader.sampler

    # One for each epoch
    assert train_sampler.set_epoch.mock_calls == [call(0), call(1)]
    # One for each epoch + sanity check
    assert val_sampler.set_epoch.mock_calls == [call(0), call(0), call(1)]

    val_dataloader = _get_dataloader()
    trainer.validate(model, val_dataloader)
    val_sampler = val_dataloader.batch_sampler.sampler if use_batch_sampler else val_dataloader.sampler

    assert val_sampler.set_epoch.mock_calls == [call(2)]
Tests that set_epoch is called on the dataloader's sampler (if any) during training and validation.
tests/tests_pytorch/loops/test_evaluation_loop.py