import torch


def unnormalize(normalized_imgs, mean, std):
    """Unnormalize image for visualization.
    
    Args:
        normalized_imgs (torch.tensor): image with shape (b, 3, h, w)
        mean (list of floats): list of ImageNet means for 3-channel RGB images
        std (list of floats): list of ImageNet standard deviations for 3-channel RGB images
    
    Returns:
        unnormalized_imgs (torch.tensor): unnormalized image with shape (b, 3, h, w)
    """
    assert(len(normalized_imgs.shape) == 4), f"Image tensor should have 4 dimensions, \
        but got tensor with {len(normalized_imgs.shape)} dimensions."
        
    # make a copy to avoid unnormalizing in place
    normalized_imgs_copy = normalized_imgs.clone()
    r, g, b = normalized_imgs_copy.split(1, dim=1)
    r.mul_(std[0]), g.mul_(std[1]), b.mul_(std[2])
    r.add_(mean[0]), g.add_(mean[1]), b.add_(mean[2])
    unnormalized_imgs = torch.cat([r, g, b], dim=1)

    return unnormalized_imgs