code: def gradient_penalty(critic, h_s, h_t):
      """Computes gradient penalty in Wasserstein distance guided representation learning
  
      Examples:
          See WDGRLTrainer and WDGRLTrainerMod in kale.pipeline.domain_adapter
      """
  
      alpha = torch.rand(h_s.size(0), 1)
      alpha = alpha.expand(h_s.size()).type_as(h_s)
      # try:
      differences = h_t - h_s
  
      interpolates = h_s + (alpha * differences)
      interpolates = torch.cat(
          (interpolates, h_s, h_t), dim=0
      ).requires_grad_()
  
      preds = critic(interpolates)
      gradients = grad(
          preds,
          interpolates,
          grad_outputs=torch.ones_like(preds),
          retain_graph=True,
          create_graph=True,
      )[0]
      gradient_norm = gradients.norm(2, dim=1)
      gradient_penalty = ((gradient_norm - 1) ** 2).mean()
      # except:
      #     gradient_penalty = 0
  
      return gradient_penalty
  
file_path: pykale/kale/predict/losses.py
imports: import torch
  import torch.nn
  from torch.autograd import grad
  from torch.nn import functional
label_desc: Use image node_id 6 for calling a global function with example usage: gradient_penalty(critic, h_s, h_t) and returns: gradient_penalty