How can I get the specified value of Tensor quickly?

If A, B are LongTensor, with the size of H*W*2, how can I make:

A[B[j,k,0],B[j,k,1]] = j,k

and when it comes to A, B are LongTensor, with the size of B*H*W*2, how can I achieve the same operation quickly?

A[0,B[0,j,k,0],B[0,j,k,1]] = j,k
A[1,B[1,j,k,0],B[1,j,k,1]] = j,k
...
A[B,B[1,j,k,0],B[B,j,k,1]] = j,k

You are confusing me, what’s B exactly? What’s the shape of A and B? Why C is here?

Not sure whether it’s what you need

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

Sorry, I didn’t describe my question clearly. I have changed it, but I think you got my point. Do you have any idea to achieve it quickly? Thanks a lot.:grin:


import torch
A=torch.ones(2,3,4,2).long()
B=torch.ones(2,3,4,2).long()    

N,H,W,_ = A.shape
b = torch.arange(N).view(N,1,1) #N,1,1
j = torch.arange(H)
k = torch.arange(W)
jk = torch.meshgrid(j,k)
jk = torch.stack(jk,2).view(1,H,W,2) #1,H,W,2
A[b,B[...,0],B[...,1]] = jk # broadcast to N,H,W,2
1 Like

Wow, this is exactly what I want. Thank you vary much!!!