Hi. I have a model where I will need to use one element from a matrix of elements. I computed a softmax probability distribution over this matrix. Now, i would like to extract the element that corresponds to the maximum probability from the softmax, and then use that element further on:
_, idx = torch.max(softmax)
desired_element = matrix[idx]

and then the desired_element is used for further operations in the model.
Is that operation differentiable?
Thanks!