A differentiable argmax function in PyTorch, for your gradient descending convenience.
A differentiable argmax function for PyTorch