How to select elements from a tensor according to a 2d index?

[07hyx06]

Suppose A is a 2D tensor with shape (3,4). B is a 2D tensor with shape (k,2) specify k selected index. The output is a tensor C with shape (k) described as:

C = []
for i in range(len(B)):
    C.append(A[B[i][0]][B[i][1]])
C = torch.tensor(C)

I wonder if there exists some Pytorch API equivalent to the above code segment?

[KFrank]

Hi 07hyx06!

@07hyx06

C = []
for i in range(len(B)):
    C.append(A[B[i][0]][B[i][1]])
C = torch.tensor(C)

I wonder if there exists some Pytorch API equivalent to the above code segment?

Yes:

A[B[:, 0], B[:, 1]]

Best.

K. Frank