That looks like a sensible approach but you need the inverse permutation to undo the sorting (you have that, sorry for the confusion). If you use input_sizes as first and an arange in the second indexing object, you can merge the loop and stack into an advanced indexing operation (there might also be an index_something function that does that for you without needing the arange).
Best regards
Thomas