# Index into tensor with indices from another?

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!

This should work:

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

Thank you! This works!