from torch.autograd import grad
preds = Critic(interpolates)
    gradients = grad(outputs=preds, inputs=interpolates,
                     grad_outputs=torch.ones_like(preds).cuda(),
                     retain_graph=True, create_graph=True, only_inputs=True)[0]
    gradients = gradients.view(gradients.size(0),  -1)
    gradient_norm = gradients.norm(2, dim=1)
    gradient_penalty = ((gradient_norm - 1)**2).mean()