Index into tensor with indices from another?

1

Hello,

I’ve been struggling for hours with this.

Say I have a 2d tensor:
A = torch.randn(10,4)

And also a 1d LongTensor of with A.shape[0] elements (e.g. 10)
I = torch.randint(0,4,(10))

I’d like an easier but equivalent way of doing this:
torch.take(A,I+torch.arange(0,I.shape[0])*A.shape[1])

That is, I want to obtain a 1-d tensor like this: from the 0th row of A, the I[0]-th value, from the first row of A the I[1]-th value, etc.

If I do A[I] it gives me a matrix (I want a 1d tensor) that selects the corresponding rows.
I also tried index_select, but I’m also getting a matrix.

Thank you!

2

This should work:

A[torch.arange(A.size(0)), I]
2 Likes
3

Thank you! This works!