code: def test_transformations(tmpdir):
      transform = TestInputTransform()
      datamodule = DataModule(
          TestInput(RunningStage.TRAINING, [1]),
          TestInput(RunningStage.VALIDATING, [1]),
          TestInput(RunningStage.TESTING, [1]),
          transform=transform,
          batch_size=2,
          num_workers=0,
      )
  
      assert datamodule.train_dataloader().dataset[0] == (0, 1, 2, 3)
      batch = next(iter(datamodule.train_dataloader()))
      assert torch.equal(
          batch, torch.tensor([[0, 1, 2, 3, 5], [0, 1, 2, 3, 5]])
      )
  
      assert datamodule.val_dataloader().dataset[0] == {"a": 0, "b": 1}
      assert datamodule.val_dataloader().dataset[1] == {"a": 1, "b": 2}
      batch = next(iter(datamodule.val_dataloader()))
  
      datamodule = DataModule(
          TestInput(RunningStage.TRAINING, [1]),
          TestInput(RunningStage.VALIDATING, [1]),
          TestInput(RunningStage.TESTING, [1]),
          transform=TestInputTransform2,
          batch_size=2,
          num_workers=0,
      )
      batch = next(iter(datamodule.val_dataloader()))
      assert torch.equal(batch["a"], torch.tensor([0, 1]))
      assert torch.equal(batch["b"], torch.tensor([1, 2]))
  
      model = CustomModel()
      trainer = Trainer(
          max_epochs=1,
          limit_train_batches=2,
          limit_val_batches=1,
          limit_test_batches=2,
          limit_predict_batches=2,
          num_sanity_val_steps=1,
      )
      trainer.fit(model, datamodule=datamodule)
      trainer.test(model, datamodule=datamodule)
  
      assert (
          datamodule.input_transform.train_per_sample_transform_called
      )
      assert datamodule.input_transform.train_collate_called
      assert (
          datamodule.input_transform.train_per_batch_transform_on_device_called
      )
      assert (
          datamodule.input_transform.train_per_sample_transform_called
      )
      assert datamodule.input_transform.val_collate_called
      assert (
          datamodule.input_transform.val_per_batch_transform_on_device_called
      )
      assert datamodule.input_transform.test_per_sample_transform_called
  
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 2 for calling a global function with example usage: test_transformations(tmpdir) without return types