import torch


def calculate_gradient_penalty(model,step,alpha, real_images, fake_images, lambd ,device):
    """Calculates the gradient penalty loss for WGAN GP"""
    eps = torch.rand(real_images.size(0), 1, 1, 1, device=device)
    eps = eps.expand_as(real_images)
    x_hat = eps * real_images + (1 - eps) * fake_images.detach()
    x_hat.requires_grad = True
    px_hat = model(x_hat,step=step,alpha=alpha)
    grad = torch.autograd.grad(outputs = px_hat.sum(),
                                    inputs = x_hat, 
                                    create_graph=True
                                    )[0]
    grad_norm = grad.view(real_images.size(0), -1).norm(2, dim=1)
    gradient_penalty = lambd * ((grad_norm  - 1)**2).mean()
    return gradient_penalty