Solved by theycallmefm in post #2 
              
              
                I think I have found a way. I am posting it here in case somebody needs it. 
bidx = torch.arange(B).view(B,1,1)
bidx = bidx.expand_as(idx).contiguous()
nidx = torch.arange(N).view(1,N,1)
nidx = nidx.expand_as(idx).contiguous()
x = x[bidx.view(-1), nidx.view(-1), idx.view(-1)].view(B,N,D)