code: def backward(ctx, grad_output):
      
      batchSize = grad_output.shape[0]
      indices, = ctx.saved_tensors
  
      g = grad_output.reshape((batchSize, -1))
      grad_x = torch.zeros_like(g, requires_grad=False)
      grad_x.scatter_(1, indices, g.gather(1, indices))
      grad_x = grad_x.reshape(grad_output.shape)
  
      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.