I have a tensor X of shape (a, b, c) and a matrix of permutation (not a permutation matrix) P of shape (a,b), where each row of P is an output of torch.randperm(). I want to shuffle X in the following way:
for i in range(a):
    Y[i] = X[i][P[i]]
return Y

what’s the best way to achieve this? Thanks.