How index from one dimension of a multi dimensional tensor using another tensor as indexes?

1

Hi! I have two arrays of shape:

items.shape=(200,100,3)
indexes.shape=(200)

I would like to pick, for all 200 items in the first dimension, the item designed by the respective item at indexes. I would like to do a more elegant solution than

torch.stack([items[i,indexes[i]] for i in range(len(indexes))])

Thanks in advance!

2

This should work:

# setup
items = torch.randn(200, 100, 3)
indexes = torch.randint(0, 100, (200,))

# reference
ref = torch.stack([items[i,indexes[i]] for i in range(len(indexes))])

# direct indexing
out = items[torch.arange(items.size(0)), indexes]

# test
print((out == ref).all())
# tensor(True)