code: def permute_target(sample: Dict[str, Any]) -> Dict[str, Any]:
      if DataKeys.TARGET in sample:
          target = sample[DataKeys.TARGET]
          if target.ndim == 2:
              target = target[None, :, :]
          sample[DataKeys.TARGET] = target.transpose((1, 2, 0))
      return sample
  
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 2 for calling a global function with example usage: permute_target(sample) and returns: sample