From b5fe5bb5259468ae3a14544a4cf6b0cc0055673d Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 6 Feb 2026 14:28:53 -0500 Subject: [PATCH] Add max_sequences_per_bin to second _PackedBatch --- grain/_src/python/experimental/example_packing/packing.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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)