text: import torch
  from torch.autograd.gradcheck import zero_gradients
  
  
  def jacobian(inputs, output):
      """
      :param inputs: Batch X Size (e.g. Depth X Width X Height)
      :param output: Batch X Classes
      :return: jacobian: Batch X Classes X Size
      """
      assert inputs.requires_grad
  
      num_classes = output.size()[1]
  
      J = torch.zeros(num_classes, *inputs.size())
      grad_output = torch.zeros(*output.size())
      if inputs.is_cuda:
          grad_output = grad_output.cuda()
          J = J.cuda()
  
      for i in range(num_classes):
          zero_gradients(inputs)
          grad_output.zero_()
          grad_output[:, i] = 1
          output.backward(grad_output, retain_graph=True)
          J[i] = inputs.grad.data
  
      return torch.transpose(J, dim0=0, dim1=1)
  