code: def backward(ctx, grad_output):
      
      indices, = ctx.saved_tensors
      grad_x = torch.zeros_like(grad_output, requires_grad=True)
  
      # Probably a better way to do it, but this is not terrible as it only loops
      # over the batch size.
      for i in range(grad_output.size(0)):
        grad_x[i, indices[i]] = grad_output[i, indices[i]]
  
      return grad_x, None, None, None
docs: In the backward pass, we set the gradient to 1 for the winning units, and 0
      for the others.