Of course. Typically, your predictions come back in the order you provide them to the predict() function. So you can just merge the data back together again at the end
The exception is torchtext which sorts by the field lengths or something. For that you need to set batch size to 1 when evaluating test results so you can maintain the same order