code: 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)]
  
imports: from unittest import mock
  from unittest.mock import Mock, call
  import pytest
  import torch
  from lightning.fabric.fabric import Fabric
  from lightning.fabric.plugins import Precision
  from lightning.fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_1
  from lightning.fabric.wrappers import _FabricDataLoader, _FabricModule, _FabricOptimizer, _unwrap_objects, is_wrapped
  from torch.utils.data import BatchSampler, DistributedSampler
  from torch.utils.data.dataloader import DataLoader
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 10 for calling a global function with example usage: test_fabric_dataloader_distributed_sampler_set_epoch(use_batch_sampler) without return types