def gradient_penalty(self, data_t, data_k, t, k):
    batch_size = data_k.size()[0]
    theta = torch.rand(batch_size, 1, 1, 1)
    theta = theta.expand_as(data_t)
    theta = theta.to(self.device)
    interpolated = ((theta * data_t) + ((1 - theta) * data_k))
    interpolated = Variable(interpolated, requires_grad=True)
    interpolated = interpolated.to(self.device)
    features_intrpltd = self.FE(interpolated)
    features_intrpltd = features_intrpltd.view(features_intrpltd.size(0), (- 1))
    prob_interpolated = self.discrm['{}{}'.format(t, k)](features_intrpltd).squeeze()
    gradients = torch_grad(outputs=prob_interpolated, inputs=interpolated, grad_outputs=torch.ones(prob_interpolated.size()).to(self.device), create_graph=True, retain_graph=True)[0]
    gradients = gradients.view(batch_size, (- 1))
    gradients_norm = torch.sqrt((torch.sum((gradients ** 2), dim=1) + 1e-12))
    return ((gradients_norm - 1) ** 2).mean()
def gradient_penalty(self, data_t, data_k, t, k):
    batch_size = data_k.size()[0]
    theta = torch.rand(batch_size, 1, 1, 1)
    theta = theta.expand_as(data_t)
    theta = theta.to(self.device)
    interpolated = ((theta * data_t) + ((1 - theta) * data_k))
    interpolated = Variable(interpolated, requires_grad=True)
    interpolated = interpolated.to(self.device)
    features_intrpltd = self.FE(interpolated)
    features_intrpltd = features_intrpltd.view(features_intrpltd.size(0), (- 1))
    prob_interpolated = self.discrm['{}{}'.format(t, k)](features_intrpltd).squeeze()
    gradients = torch_grad(outputs=prob_interpolated, inputs=interpolated, grad_outputs=torch.ones(prob_interpolated.size()).to(self.device), create_graph=True, retain_graph=True)[0]
    gradients = gradients.view(batch_size, (- 1))
    gradients_norm = torch.sqrt((torch.sum((gradients ** 2), dim=1) + 1e-12))
    return ((gradients_norm - 1) ** 2).mean()