import torch


def np2torch(array, dtype=None):
    """
    Convert a numpy array to torch tensor convention.
    If 4D -> [b, h, w, c] to [b, c, h, w]
    If 3D -> [h, w, c] to [c, h, w]

    :param array: Numpy array
    :param dtype: Target tensor dtype
    :return: Torch tensor
    """
    d = array.ndim
    perm = [0, 3, 1, 2] if d == 4 else \
           [2, 0, 1] if d == 3 else \
           [0, 1]

    tensor = torch.from_numpy(array).permute(perm)
    return tensor.type(dtype) if dtype else tensor.float()