Solved by Eta_C in post #2 
              
              
                Try 
import torch 

x = torch.rand(3, 5)
indices = torch.argsort(torch.rand(*x.shape), dim=-1)
result = x[torch.arange(x.shape[0]).unsqueeze(-1), indices]