Solved by tom in post #2 
              
              
                It seems rather bold to say that it’s nearly impossible to do without for loops. 
For example, I can replace the line defining inverse_permutation_indexes with 
inverse_permutation_indexes = torch.empty_like(permutation_indexes).scatter_(dim=1, index=permutation_indexes, src=y)

and get the computat…