code: def remove_extra_dimensions(batch: Dict[str, Any]):
      if isinstance(batch[DataKeys.INPUT], list):
          assert len(batch[DataKeys.INPUT]) == 1
          batch[DataKeys.INPUT] = batch[DataKeys.INPUT][0]
      return batch
  
file_path: lightning-flash/src/flash/image/segmentation/input_transform.py
imports: from dataclasses import dataclass
  from typing import Any, Callable, Dict, Tuple
  from flash.core.data.io.input import DataKeys
  from flash.core.data.io.input_transform import InputTransform
  from flash.core.data.transforms import AlbumentationsAdapter, ApplyToKeys
  from flash.core.utilities.imports import _ALBUMENTATIONS_AVAILABLE, _TORCHVISION_AVAILABLE, requires
label_desc: Use image node_id 3 for calling a global function with example usage: remove_extra_dimensions(batch) and returns: batch