@aron-bordin So using this, if you have multiple mini-batches (with different sorting permutations) going in different RNNs,  this forward pass returns them in the original order and then they are directly comparable. Pytorch takes care of this permutation while backpropagating. I hope I understand this correctly.
One question, how did you check that the grads are permuted correctly?
Thanks!