code: 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)]
  
imports: from unittest import mock
  from unittest.mock import Mock, call
  import pytest
  import torch
  from lightning.fabric.accelerators.cuda import _clear_cuda_memory
  from lightning.pytorch import LightningModule, Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
  from lightning.pytorch.utilities import CombinedLoader
  from torch.utils.data.dataloader import DataLoader
  from torch.utils.data.sampler import BatchSampler, RandomSampler
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 2 for calling a global function with example usage: test_evaluation_loop_sampler_set_epoch_called(tmp_path, use_batch_sampler) and returns: DataLoader, DataLoader