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 elements (e.g. 10)
I = torch.randint(0,4,(10))
I’d like an easier but equivalent way of doing this:
That is, I want to obtain a 1-d tensor like this: from the 0th row of A, the
I-th value, from the first row of A the
I-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.