Index tensor must have same dimensions as input tensor

1

Help me :slight_smile:

import torch
import random
t = torch.randn((32,100))
a = torch.tensor([random.randint(0,100) for a in range(32)])
torch.gather(t, 1, a)

Traceback (most recent call last):
File “”, line 6, in
RuntimeError: invalid argument 4: Index tensor must have same dimensions as input tensor at C:\w\1\s\windows\pytorch\aten\src\TH/generic/THTensorEvenMoreMath.cpp:453

2

If you want to use a as the index in dim1 for t, this code should work:

t[torch.arange(t.size(0)), a]
3

torch.gather(t, 1, a.reshape(32, 1))

I had a mistake. It also works!!