code: def test_dataloaders_with_sampler(mock_dataloader, sampler, callable):
      train_input = TestInput(RunningStage.TRAINING, [1])
      datamodule = DataModule(
          train_input,
          TestInput(RunningStage.VALIDATING, [1]),
          TestInput(RunningStage.TESTING, [1]),
          batch_size=2,
          num_workers=0,
          sampler=sampler,
      )
  
      assert datamodule.sampler is sampler
      dl = datamodule.train_dataloader()
  
      if callable:
          sampler.assert_called_once_with(train_input)
  
      kwargs = mock_dataloader.call_args[1]
      assert "sampler" in kwargs
      assert kwargs["sampler"] is (
          sampler.return_value if callable else sampler
      )
      for dl in [
          datamodule.val_dataloader(),
          datamodule.test_dataloader(),
      ]:
          kwargs = mock_dataloader.call_args[1]
          assert "sampler" not in kwargs
  
imports: from dataclasses import dataclass
  from typing import Callable, Dict
  from unittest.mock import MagicMock, NonCallableMock, patch
  import numpy
  import pytest
  import torch
  from flash import Task, Trainer
  from flash.core.data.data_module import DataModule, DatasetInput
  from flash.core.data.io.input import Input
  from flash.core.data.io.input_transform import InputTransform
  from flash.core.utilities.imports import _TOPIC_CORE_AVAILABLE, _TOPIC_IMAGE_AVAILABLE, _TORCHVISION_AVAILABLE
  from flash.core.utilities.stages import RunningStage
  from pytorch_lightning import seed_everything
  from torch.utils.data import Dataset
  from tests.helpers.boring_model import BoringModel
label_desc: Use image node_id 4 for calling a global function with example usage: test_dataloaders_with_sampler(mock_dataloader, sampler, callable) without return types