import torch


def np2torch(nparray):
    """Convert numpy array to torch tensor
       For array with more than 3 channels, it is better to use an input array in the format BxHxWxC

       Args:
           numpy array (ndarray) BxHxWxC
       Returns:
           torch tensor (tensor) BxCxHxW"""

    tensor = torch.Tensor(nparray)

    if tensor.ndim == 2:
        return tensor
    if tensor.ndim == 3:
        height, width, channels = tensor.shape
        if channels <= 3:  # Single image with more channels (HxWxC)
            return tensor.permute(2, 0, 1)

    if tensor.ndim == 4:  # More images with more channels (BxHxWxC)
        return tensor.permute(0, 3, 1, 2)

    return tensor