Jordi_de_la_Torre:

          # restore initial ordering                                              
          last1 = last1[initial_idx1]                                             
          last2 = last2[initial_idx2]                                             



The above code doesn’t restore initial ordering. Use torch.Tensor.scatter_ instead.