import torch


def get_gradient(crit, real, fake, epsilon):
    """
    Returns the gradient of the critic's scores with respect to mixes of real and fake images
    :param crit: a critic model
    :param real: a batch of real_images
    :param fake: a batch of fake_images
    :param epsilon: a vector of the uniformly random proportions of real and fake per mixed image
    :return: gradient: the gradient of the critic's scores wrt the mixed image
    """
    # Mix the images together
    mixed_images = real * epsilon + fake * (1 - epsilon)
    mixed_scores = crit(mixed_images)
    gradient = torch.autograd.grad(
        inputs=mixed_images,
        outputs=mixed_scores,
        grad_outputs=torch.ones_like(mixed_scores),
        create_graph=True,
        retain_graph=True
    )[0]
    return gradient