def METHOD_NAME(self, order):
    # Creates an index that undoes a sort: xs==xs[order][inverse_sort(order)]
    return torch.empty_like(order).scatter_(0, order, torch.arange(0, order.size(0), device=order.device))