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!
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])
Yes, actually [0, 4] would be correct! Thank you so much for the fast reply