How to get the value of tensor

A: shape [b,n,h,w,c]
B: shape [b,n,h,w,2]
How can i get the value of A[bi,ni,B[bi,ni,hi,wi,0],B[bi,ni,hi,wi,1],:] ?

If you want to pick
1 - just one value, A[...].item() should work.
2 - an array A[.., :].numpy() should work.

thank you for your reply. But I do need to create a new tensor C which C[bi,ni,hi,wi,:]=A[bi,ni,B[bi,ni,hi,wi,0],B[bi,ni,hi,wi,1],:] . In other word , I want to use 2d_index (hi,wi) to get the value of A[…,hi,wi,…] and generate a new tensor C , and C.shape=A.shape .

I am confused about why just writing C[bi,ni,hi,wi,:]=A[bi,ni,B[bi,ni,hi,wi,0],B[bi,ni,hi,wi,1],:] is not working?

Yeah , It can work. but this method need many loop functions , which may work slow . i want to know if there are any function in pytorch or other libraries, which can work quickly. .

I find the function . Firstly, reshape the A and B into b,n,h*w,: and then use torch.gather function could solve this problem. thanks.

1 Like