Say I have a matrix X which is of size (x1, x2, x3, x4, x5) and I do,
values, locations = torch.max(X, 2)

Is there a good way of being able to do autograd for the locations, rather than the values? Yes, I can multiply the values by the locations to come up with a weighted metric, but what I’m really looking for is just the locations.
Right now after I do backward(), I get None for my gradient.