def backward(ctx, grad_output):
    """
    In the backward pass, we set the gradient to 1 for the winning units, and 0
    for the others.
    """
    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
In the backward pass, we set the gradient to 1 for the winning units, and 0
    for the others.