Randomly get index of one of the maximum values in a PyTorch tensor
I need to perform something similar to the built-in torch.argmax() function on a one-dimensional tensor, but instead of picking the index of the first of the maximum values, I want to be able to pick a random index of one of the maximum values. For example:
my_tensor = torch.tensor([0.1, 0.2, 0.2, 0.1, 0.1, 0.2, 0.1])
index_1 = random_max_val_index_fn(my_tensor)
index_2 = random_max_val_index_fn(my_tensor)

print(f&quot;{index_1}, {index_2}&quot;)
&gt; 5, 1


You can get the indexes of all the maximums first and then choose randomly from them:
def rand_argmax(tens):
    max_inds, = torch.where(tens == tens.max())
    return np.random.choice(max_inds)

sample runs:
&gt;&gt;&gt; my_tensor = torch.tensor([0.1, 0.2, 0.2, 0.1, 0.1, 0.2, 0.1])
&gt;&gt;&gt; rand_argmax(my_tensor)
2
&gt;&gt;&gt; rand_argmax(my_tensor)
5
&gt;&gt;&gt; rand_argmax(my_tensor)
2
&gt;&gt;&gt; rand_argmax(my_tensor)
1

