What's the best practice to make a tensor A that A[i]=B[C[i]]

1

I got tensor C (of size N ) and tensor B (of size M), and I want to make a tensor A which
A[i]=B[C[i]],
Actually I can use A = [B[x] for x in C], which create a list(of size N) of tensor(scalar) but not a tensor(of size N), and for-loop seems slow without vectorization.
I think there must have a better way to do it, but i’m new to torch, so i’m looking for help.
Thanks in advance.

2

Direct indexing should work:

N, M = 4, 5

C = torch.randint(0, M, (N,))
B = torch.randn(M)

A_slow = torch.stack([B[x] for x in C])

A = B[C[torch.arange(C.size(0))]]

print((A_slow == A).all())
# tensor(True)
1 Like
3

That works, Thanks a lot!