def test_fabric_dataloader_distributed_sampler_set_epoch(use_batch_sampler):
    """Test that the FabricDataLoader calls `set_epoch()` on the wrapped sampler if applicable."""
    dataset = range(3)
    sampler = DistributedSampler(dataset, num_replicas=2, rank=0)
    sampler.set_epoch = Mock()

    if not use_batch_sampler:
        dataloader = DataLoader(dataset, sampler=sampler)
    else:
        batch_sampler = BatchSampler(sampler, batch_size=1, drop_last=False)
        dataloader = DataLoader(dataset, batch_sampler=batch_sampler)

    fabric_dataloader = _FabricDataLoader(dataloader)
    iterator_epoch_0 = iter(fabric_dataloader)
    sampler.set_epoch.assert_not_called()

    next(iterator_epoch_0)
    # .set_epoch() gets called before the first sample gets fetched from the wrapped dataloader
    assert sampler.set_epoch.mock_calls == [call(0)]

    next(iterator_epoch_0)
    assert sampler.set_epoch.mock_calls == [call(0)]

    iterator_epoch_1 = iter(fabric_dataloader)
    assert sampler.set_epoch.mock_calls == [call(0)]

    next(iterator_epoch_1)
    # with every new iterator call, the epoch increases
    assert sampler.set_epoch.mock_calls == [call(0), call(1)]
Test that the FabricDataLoader calls `set_epoch()` on the wrapped sampler if applicable.
tests/tests_fabric/test_wrappers.py