Support generator in BucketBatchSampler

PetrochukM · Jan 26, 2021 · 0d9824e · 0d9824e
1 parent be1a08e
commit 0d9824e
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/torchnlp/samplers/bucket_batch_sampler.py b/torchnlp/samplers/bucket_batch_sampler.py
@@ -54,10 +54,11 @@ def __init__(self,
                  sort_key=identity,
                  bucket_size_multiplier=100):
         super().__init__(sampler, batch_size, drop_last)
-        self.sort_key = sort_key
-        self.bucket_sampler = BatchSampler(sampler,
-                                           min(batch_size * bucket_size_multiplier, len(sampler)),
-                                           False)
+        self.sort_key = sort_key       
+        _bucket_size = batch_size * bucket_size_multiplier
+        if hasattr(sampler, "__len__"):
+            _bucket_size = min(_bucket_size, len(sampler))
+        self.bucket_sampler = BatchSampler(sampler, _bucket_size, False)
 
     def __iter__(self):
         for bucket in self.bucket_sampler: