diff --git a/grain/_src/python/experimental/example_packing/packing.py b/grain/_src/python/experimental/example_packing/packing.py index a1a16463c..6f2df2e85 100644 --- a/grain/_src/python/experimental/example_packing/packing.py +++ b/grain/_src/python/experimental/example_packing/packing.py @@ -232,7 +232,10 @@ def __call__( if not element_added_to_batch: yield self._cur_batch.get_packed_batch() # Main yield self._cur_batch = _PackedBatch( - element.data, self.batch_size, self.length_struct + element.data, + self.batch_size, + self.length_struct, + self.max_sequences_per_bin, ) self._cur_batch.try_add_to_batch(element)