import torch


def mask_to_tensor(mask):
    """
    transforms (H, W) or (H, W, Class) numpy array into (Class, H, W) tensor
    """
    mask = torch.from_numpy(mask).float()
    if mask.ndimension() == 2:
        mask = mask.unsqueeze(-1)
    return mask.permute(2, 0, 1)