Chose random element from tensor excluding certain index

I am wondering what the most efficient way is, to randomly select an index i of a tensor when, say, every index except i=1 is possible.

a = torch.tensor([0,1,2,3,4])
# randomly sample from a one index but not the 1 at position 1

I want to avoid any loops, since I need to do this in my getitem function

torch.multinomial might work for you:

a = torch.arange(5).float()
weight = torch.ones_like(a)
weight[1] = 0.
a[torch.multinomial(weight, num_samples=1)]
2 Likes

thank you! your advice is always so helpful =)