In a data parallel setup the output of the models will be gathered in the same way as the data was chunked in before the forward pass.
I.e. if you increase the batch size and make sure that the order of the samples inside the batch is correct, the extracted features will have the same order.
Freezing all parameters should also work, so you wouldn’t need the no_grad statement.