fix: Respect max batch size in PipelineTrainer

vivekkalyan · vivekkalyan · commit d7a9a06ef1db · 2026-03-20T18:07:39.000-07:00
diff --git a/src/art/pipeline_trainer/trainer.py b/src/art/pipeline_trainer/trainer.py
@@ -552,7 +552,7 @@ async def _collect_batch(
                 continue
             batch.append(item)
 
-        while not saw_sentinel:
+        while not saw_sentinel and len(batch) < self.max_batch_size:
             try:
                 item = self._output_queue.get_nowait()
             except asyncio.QueueEmpty:
