Slice matrix with another matrix

1

Hi,
I would like the following without using for loop:

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

for i in range(len(a)):
    c = a[i][b[i]]
    print(c)

a - matrix of values,
b - matrix of indices

Any ideas?

2

All right, found!

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

c=a[torch.arange(a.size(0)).unsqueeze(1),b]