Is there a method to make index having gradient function?
i.e.

import torch
h = torch.randn(1,2,5, requires_grad=True); print(h)
val,idx = h.max(1, keepdim=True)
print(val)
print(idx)
print(val)

outputs are:

tensor([[[-0.5372, -0.4683,  0.4891, -0.1686, -0.4147],
[-1.4412,  1.2837, -0.4467,  0.1731,  1.3256]]], requires_grad=True)
tensor([[[-0.5372,  1.2837,  0.4891,  0.1731,  1.3256]]],
grad_fn=)
tensor([[[0, 1, 0, 1, 1]]])

I want the tensor([[[0, 1, 0, 1, 1]]]) to have gradient function.