brisker:

how about this operation:
v = torch.autograd.Variable(torch.Tensor([0.8,0.1,0.1,0]), requires_grad=True)
out = (v==torch.max(v)).float()
does variable v and out have gradients?


You can check that out.requires_grad=False. @SimonW is right by notifying argmax (but here it’s rather an indicator function than the true “argmax”) is not differentiable. However, it’s possible to set the derivative of the function as  f'(x)=0 for all x.
For instance,
f(x,t) = Softmax(x*t) --> argmax(x) when t --> +inf
and respectivly, f'(x,t) --> 0
One simple solution could be:
v = Variable(torch.Tensor([0.8,0.1,0.1,0]), requires_grad=True)
y = F.softmax(v*100)