import torch


def numpy2torch(img):
    """
    Converts a numpy image to torch format
    """
    img = torch.from_numpy(img)
    if img.dim() == 4:
        img = img.permute(0, 3, 1, 2).contiguous()
    elif img.dim() == 3:
        img = img.permute(2, 0, 1).contiguous()
    return img