code: # Compute MMD distance using pytorch
  
  import torch
  import torch.nn as nn
  
  
  class MMD_loss(nn.Module):
      def __init__(self, kernel_type='rbf', kernel_mul=2.0, kernel_num=5):
          super(MMD_loss, self).__init__()
          self.kernel_num = kernel_num
          self.kernel_mul = kernel_mul
          self.fix_sigma = None
          self.kernel_type = kernel_type
  
      def guassian_kernel(self, source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):
          n_samples = int(source.size()[0]) + int(target.size()[0])
          total = torch.cat([source, target], dim=0)
          total0 = total.unsqueeze(0).expand(
              int(total.size(0)), int(total.size(0)), int(total.size(1)))
          total1 = total.unsqueeze(1).expand(
              int(total.size(0)), int(total.size(0)), int(total.size(1)))
          L2_distance = ((total0-total1)**2).sum(2)
          if fix_sigma:
              bandwidth = fix_sigma
          else:
              bandwidth = torch.sum(L2_distance.data) / (n_samples**2-n_samples)
          bandwidth /= kernel_mul ** (kernel_num // 2)
          bandwidth_list = [bandwidth * (kernel_mul**i)
                            for i in range(kernel_num)]
          kernel_val = [torch.exp(-L2_distance / bandwidth_temp)
                        for bandwidth_temp in bandwidth_list]
          return sum(kernel_val)
  
      def linear_mmd2(self, f_of_X, f_of_Y):
          loss = 0.0
          delta = f_of_X.float().mean(0) - f_of_Y.float().mean(0)
          loss = delta.dot(delta.T)
          return loss
  
      def forward(self, source, target):
          if self.kernel_type == 'linear':
              return self.linear_mmd2(source, target)
          elif self.kernel_type == 'rbf':
              batch_size = int(source.size()[0])
              kernels = self.guassian_kernel(
                  source, target, kernel_mul=self.kernel_mul, kernel_num=self.kernel_num, fix_sigma=self.fix_sigma)
              with torch.no_grad():
                  XX = torch.mean(kernels[:batch_size, :batch_size])
                  YY = torch.mean(kernels[batch_size:, batch_size:])
                  XY = torch.mean(kernels[:batch_size, batch_size:])
                  YX = torch.mean(kernels[batch_size:, :batch_size])
                  loss = torch.mean(XX + YY - XY - YX)
              torch.cuda.empty_cache()
              return loss
  
apis: torch.sum torch.cuda.empty_cache torch.no_grad torch.exp torch.cat torch.mean