From ced931d08d2309000af07da360f9b29fec27be1b Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 17 Feb 2026 08:27:35 -0800 Subject: [PATCH] Add missing max_sequences_per_bin for packed inputs ported from https://github.com/google/grain/pull/1215 PiperOrigin-RevId: 871330744 --- 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)