I guess using the gather makes the same mistakes as @Jordi_de_la_Torre’s code.
I think the correct way is to using the .scatter_  as suggested by @abhishek0318.
Or still using the indexing method as following:
    sorted_lengths, sorted_id = lengths.sort(0, descending=True)
    print sorted_lengths

    sorted_sorted_id, initial_id = sorted_id.sort(0, descending=False)

    sorted_input = input[sorted_id]

    print sorted_input
    print sorted_input[initial_id]

I wrote a toy program to check the correctness of this usage.
About the autograd safety, I think the permuting the tensor index is ok and won’t cause any
gradient problem.