2a1a6851344172e0134f3c5f4f5c1021975f2812,torchnlp/samplers/bucket_batch_sampler.py,BucketBatchSampler,__iter__,#BucketBatchSampler#,37

Before Change

            batch_size, drop_last)

    def __iter__(self):
        batches = list(super().__iter__())
        if self.last_batch_first:
            last_batch = batches.pop()
        if self.shuffle:
            random.shuffle(batches)

After Change

            return get_batches()
        else:
            batches = list(get_batches())
            indices = heapq.nlargest(
                5,
                range(len(batches)),
                key=lambda i: len(pickle.dumps([self.data[j] for j in batches[i]])))
            front = [batches[i] for i in indices]
            for i in sorted(indices, reverse=True):
                batches.pop(i)
            batches[0:0] = front
            return iter(batches)

Instances