Solved by KFrank in post #2 
              
              
                Hi Luiz! 

At the cost of  n log (n) time complexity, you can use argsort()* and 
then use gather() to index back into your input tensor: 
>>> import torch
>>> torch.__version__
'1.9.0'
>>> in_x = torch.tensor([[[[0.,   0.,   2.,  20., 250.,   0.,   0.,   0.,   0.,   0., 250.,
...                   …