How do I index a tensor with a tensor (not quite index_select)?

#1

Let’s say I have a tensor a = torch.tensor(([0,1,2],[3,4,5])). I have another tensor, b = torch.tensor([0,1]), that I want to use to index a, such that we obtain the 0th position of the first row, 1st position of the second row (resulting tensor is [0, 3]). How can this be done? It’s not the same as index_select, which finds the 0th and 1st positions for every row. Thanks!

#2

Wouldn’t the result be [0, 4] in this case?
If so, this should work:

a = torch.tensor(([0,1,2],[3,4,5]))
idx = torch.tensor([0, 1])
a[torch.arange(a.size(0)), idx]
# > tensor([0, 4])
1 Like
#3

Yes, actually [0, 4] would be correct! Thank you so much for the fast reply :slight_smile: