import torch


def total_variation_loss(img):
    """
    This method calculates total variation loss from given image.
    Pixel difference between current and right pixel, and between current and below pixel is used in calculation.
    Sum of squares is calculated and divided by count.
    Division is not specified in paper but without division, total variation loss is very huge and dominates total loss.
    """
    s = img.size()
    count = torch.prod(torch.tensor(s[1:]))

    slide_1 = img[:, :, 1:, :]
    r1 = (img[:, :, :-1, :] - slide_1) ** 2

    slide_2 = img[:, :, :, 1:]
    r2 = (img[:, :, :, :-1] - slide_2) ** 2

    return (r1.sum() + r2.sum()) / count