code: def test_datapipeline_transformations_overridden_by_task():
      # define input transforms
      class ImageInput(Input):
          def load_data(self, folder):
              # from folder -> return files paths
              return ["a.jpg", "b.jpg"]
  
          def load_sample(self, path):
              # from a file path, load the associated image
              return np.random.uniform(0, 1, (64, 64, 3))
  
      class ImageClassificationInputTransform(InputTransform):
          def per_sample_transform(self) -> Callable:
              return T.Compose([T.ToTensor()])
  
          def per_batch_transform_on_device(self) -> Callable:
              return T.Normalize(
                  [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
              )
  
      class OverrideInputTransform(InputTransform):
          def per_sample_transform(self) -> Callable:
              return T.Compose([T.ToTensor(), T.Resize(128)])
  
      # define task which overrides transforms using set_state
      class CustomModel(Task):
          def __init__(self):
              super().__init__(
                  model=torch.nn.Linear(1, 1),
                  loss_fn=torch.nn.MSELoss(),
              )
  
              # override default transform to resize images
              self.input_transform = OverrideInputTransform
  
          def training_step(self, batch, batch_idx):
              assert batch.shape == torch.Size([2, 3, 128, 128])
              assert torch.max(batch) <= 1.0
              assert torch.min(batch) >= 0.0
  
          def validation_step(self, batch, batch_idx):
              assert batch.shape == torch.Size([2, 3, 128, 128])
              assert torch.max(batch) <= 1.0
              assert torch.min(batch) >= 0.0
  
      transform = ImageClassificationInputTransform()
      datamodule = DataModule(
          ImageInput(RunningStage.TRAINING, [1]),
          ImageInput(RunningStage.VALIDATING, [1]),
          transform=transform,
          batch_size=2,
          num_workers=0,
      )
  
      # call trainer
      model = CustomModel()
      trainer = Trainer(
          max_epochs=1,
          limit_train_batches=2,
          limit_val_batches=1,
          num_sanity_val_steps=1,
      )
      trainer.fit(model, datamodule=datamodule)
  
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 3 for calling a global function with example usage: test_datapipeline_transformations_overridden_by_task() and returns: list, np, T, T, T