import torch


def get_gradient_penalty_from_gradient(gradient: torch.Tensor) -> torch.Tensor:
    """
    Return the gradient penalty, given a gradient.
    Given a batch of image gradients, we calculate the magnitude of each image's gradient and penalize the mean
    quadratic distance of each magnitude to 1.
    :param (torch.Tensor) gradient: the gradient of the critic's scores, with respect to the mixed image
    :return: a scalar torch.Tensor object containing the gradient penalty
    """
    # Flatten the gradients so that each row captures one image
    gradient = gradient.view(len(gradient), -1)
    # Calculate the magnitude of every row
    gradient_norm = gradient.norm(2, dim=1)
    # Penalize the mean squared distance of the gradient norms from 1
    return torch.mean((torch.ones_like(gradient_norm) - gradient_norm) ** 2)