Solved by ptrblck in post #2 
              
              
                You could use a scatter_/gather approach as seen here: 
B, N, k = 2, 3, 4

P = torch.randint(0, N, (B, N, k))
W = torch.randn(B, N, N)
A = torch.zeros(B, N, N)
A_ = A.clone()

# loop
for i in range(B):
    for ii in range(N):
        for j in range(k):
            A[i][ii][P[i][ii][j]] = W[i][ii][P[…