Hi,
You don’t actually need to create 2D indices, you can use gather for that:
import torch

x = torch.empty(10, 20).uniform_().float(); y = torch.randint(0, 20, (10,)).long()
print(x)
print(y)

out = x.gather(1, y.unsqueeze(-1))
print(out)