Fix Lint for Bucketizer (#296)

Summary: Fix mypy Error introduced by `max_token_bucketize` Pull Request resolved: #296 Reviewed By: NivekT Differential Revision: D34895420 Pulled By: ejguan fbshipit-source-id: 531210ea9e008320998231df8c3569d8ef7696fc
pytorch · Mar 15, 2022 · 8992d09 · 8992d09
1 parent a78843b
commit 8992d09
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torchdata/datapipes/iter/transform/bucketbatcher.py b/torchdata/datapipes/iter/transform/bucketbatcher.py
@@ -211,18 +211,18 @@ def __iter__(self) -> Iterator[DataChunk[T_co]]:
             if len(buffer) == self.buffer_size:
                 length, token = heapq.heappop(buffer)
                 if batch_size + length > self.max_token_count:
-                    yield batch
+                    yield DataChunk(batch)
                     batch = []
                     batch_size = 0
                 batch.append(token)
                 batch_size += length
         while buffer:
             length, token = heapq.heappop(buffer)
             if batch_size + length > self.max_token_count:
-                yield batch
+                yield DataChunk(batch)
                 batch = []
                 batch_size = 0
             batch.append(token)
             batch_size += length
         if batch:
-            yield batch
+            yield DataChunk(batch)