code: def test_overfit_batch_limits_train(overfit_batches):
      class CustomDataModule(ClassifDataModule):
          def train_dataloader(self):
              return DataLoader(
                  SklearnDataset(
                      self.x_train,
                      self.y_train,
                      self._x_type,
                      self._y_type,
                  ),
                  batch_size=self.batch_size,
                  shuffle=True,
              )
  
      model = ClassificationModel()
      dm = CustomDataModule()
  
      # original train loader which should be replaced in all methods
      train_loader = dm.train_dataloader()
      assert isinstance(train_loader.sampler, RandomSampler)
  
      # Create a reference train dataloader without shuffling.
      train_loader = DataLoader(
          dm.train_dataloader().dataset, shuffle=False
      )
      (xa, ya) = next(iter(train_loader))
      train_loader = DataLoader(
          dm.train_dataloader().dataset, shuffle=True
      )
      full_train_samples = len(train_loader)
  
      # set the model loaders
      model.train_dataloader = lambda: train_loader
  
      # test train loader applies correct limits
      trainer = Trainer(overfit_batches=overfit_batches)
      model.trainer = trainer
      trainer.strategy.connect(model)
      trainer._data_connector.attach_dataloaders(model=model)
      trainer.fit_loop.setup_data()
      expected_batches = (
          int(overfit_batches * full_train_samples)
          if isinstance(overfit_batches, float)
          else overfit_batches
      )
      assert trainer.num_training_batches == expected_batches
  
      # make sure the loaders are the same
      (xb, yb) = next(iter(trainer.train_dataloader))
      assert torch.eq(xa, xb).all()
      assert torch.eq(ya, yb).all()
  
file_path: pytorch-lightning/tests/tests_pytorch/trainer/flags/test_overfit_batches.py
imports: import os
  from unittest import mock
  import pytest
  import torch
  from lightning.pytorch import Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
  from lightning.pytorch.trainer.states import RunningStage
  from torch.utils.data import DataLoader, DistributedSampler, RandomSampler, Sampler, SequentialSampler
  from tests_pytorch.helpers.datamodules import ClassifDataModule
  from tests_pytorch.helpers.datasets import SklearnDataset
  from tests_pytorch.helpers.runif import RunIf
  from tests_pytorch.helpers.simple_models import ClassificationModel
label_desc: Use image node_id 4 for calling a global function with example usage: test_overfit_batch_limits_train(overfit_batches) and returns: DataLoader