Indexing a 1-dimensional tensor with another 2-dimensional tensor

1

Hi there

I am trying to index a one-dimensional tensor with a two-dimensional index.

In numpy, the following works:

a = np.array([5 6 7 8 9])
idx = 
[[1 2]
 [4 4]
 [3 3]
 [0 1]
 [2 0]]

# such that
a[idx] = 
[[6 7]
 [9 9]
 [8 8]
 [5 6]
 [7 5]]

How do I achieve the same thing in pytorch 1.0?

Doing

a = torch.tensor(a)
idx = torch.tensor(idx)
a[idx]

yields:

IndexError: too many indices for tensor of dimension 1
2

Your code just works fine:

a = torch.tensor([5, 6, 7, 8, 9])
idx = torch.tensor([[1, 2],
                    [4, 4],
                    [3, 3],
                    [0, 1],
                    [2, 0]])

a[idx]
> tensor([[6, 7],
        [9, 9],
        [8, 8],
        [5, 6],
        [7, 5]])
3

Uh yeah, forgot to convert the idx to a torch.tensor in my actual code. It does not seem to work if idx is a numpy array, though. Anyway, thank you :slight_smile: