import torch


def tv_loss(img, tv_weight=1E-6):
    """
    Compute total variation loss.

    Inputs:
    - img: PyTorch Variable of shape (1, 3, H, W) holding an input image.
    - tv_weight: Scalar giving the weight w_t to use for the TV loss.

    Returns:
    - loss: PyTorch Variable holding a scalar giving the total variation loss
      for img weighted by tv_weight.
    """
    N, C, H, W = img.size()
    f = img[:, :, :-1, :-1]
    g = img[:, :, :-1, 1:]
    h = img[:, :, 1:, :-1]
    return tv_weight * torch.sum((f - g) ** 2. + (f - h) ** 2.)