From dc7fff5bf758d8cf99a4318c676eb99f1ebab53b Mon Sep 17 00:00:00 2001 From: Jake Poznanski Date: Sun, 29 Jun 2025 19:52:53 +0000 Subject: [PATCH] Collator fix --- olmocr/train/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/olmocr/train/train.py b/olmocr/train/train.py index 3eb7ed9..8bde005 100644 --- a/olmocr/train/train.py +++ b/olmocr/train/train.py @@ -68,7 +68,7 @@ class QwenDataCollator: 'input_ids': torch.stack(batch['input_ids']), 'attention_mask': torch.stack(batch['attention_mask']), 'labels': torch.stack(batch['labels']), - 'pixel_values': batch['pixel_values'], # Keep as list for now + 'pixel_values': torch.stack(batch['pixel_values']), # Stack into tensor 'image_grid_thw': torch.stack(batch['image_grid_thw']) }