I have tensor pos which is (N,M,L,3), a second tensor env_scans (which is a grid) of size (N,K,3), and a third tensor indices (N,M,L) which contains indices (from 0 to K-1). I want to do the following:
for n in range(N):
for m in range(M):
for l in range(L):
pos[n,m,l, 2] = env_scans[n, indices[n, m, l], 2]
Is there a way to vectorise this to avoid the for loops? So far I managed to vectorise across the batch dimension (N):
for m in range(M):
for l in range(L):
pos[:,m, l, 2] = env_scans[torch.arange(N),indices[:,m, l], 2]
and this works, but I’m not sure how to expand it further for the remaining dimensions as well. I’d appreciate any clues/help!