Selecting entries of a pytorch tensor with another tensor
I have a tensor a with with float entries and torch.Size([64,2]) and I also have a tensor b with torch.Size([64]). The entries of b are only 0 or 1.
I would like to get a new tensor c with torch.Size([64]) such that c[i] == a[i,b[i]] for every index i. How can I do that?
My attempt
I tried with torch.gather but without success. The following code gives me RuntimeError: Index tensor must have the same number of dimensions as input tensor
import torch
a = torch.zeros([64,2])
b = torch.ones(64).long()
torch.gather(input=a, dim=1,index=b)

Any help will be highly appreciated!

You can perform this straight with an indexing of a on both dimensions:

On dimension=0: a &quot;sequential&quot; indexing using torch.arange.

On dimension=1: indexing using b.


All in all, this gives:
&gt;&gt;&gt; a[torch.arange(len(a)), b]


Alternatively you can use torch.gather, the operation you are looking for is:
# c[i] == a[i,b[i]]

The provided gather operation when applied on dim=1 provides something like:
# c[i,j] == a[i,b[i,j]]

As you can see, we need to account for the difference in shapes between between a and b. To do so, you can unsqueeze a singleton dimension on b (annotated with the letter j above) such than #b=(64, 1), for instance with b.unsqueeze(-1) or b[...,None]:
&gt;&gt;&gt; a.gather(dim=1, index=b[...,None]).flatten()

