import torch


def tv_loss(img, alpha=2):
    """
    Compute total variation loss.

    Args:
    -- img: PyTorch Variable of shape (N, 3, H, W) holding an input image.
    -- alpha: alpha norm.

    Returns:
    -- loss: PyTorch Variable holding a scalar giving the total variation loss for img.
    """
    N = img.shape[0]
    w_variance = torch.sum(torch.pow(img[:, :, :, :-1] - img[:, :, :, 1:], alpha))
    h_variance = torch.sum(torch.pow(img[:, :, :-1, :] - img[:, :, 1:, :], alpha))
    loss = (h_variance + w_variance) / N
    return loss