def sort_predictions(data_loader, results):
    """Recovers the dataset order when bucketing samplers are used."""
    if getattr(data_loader.batch_sampler, 'store_indices', False):
        results = [results[i] for i, j in sorted(
            enumerate(data_loader.batch_sampler.orig_idxs), key=lambda k: k[1])]
    return results