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
Convert the target mask to long and remove the channel dimension.
src/flash/image/segmentation/input_transform.py