for i in ...: for j in ...: for k in ...: outputs[i][j][k] = inputs[indices[i][j]][indices[i][k]]
is there any indexing method can help me do this?
i is the batch dim
I solved this by
inputs[indices] and then applying