code: def prepare_target(batch: Dict[str, Any]) -> Dict[str, Any]:
      """Convert the target mask to long and remove the channel dimension."""
      if DataKeys.TARGET in batch:
          batch[DataKeys.TARGET] = (
              batch[DataKeys.TARGET].squeeze().long()
          )
      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 1 for calling a global function with example usage: prepare_target(batch) and returns: batch