yujunjun/olmocr
1
0
Fork 0
mirror of https://github.com/allenai/olmocr.git synced 2025-10-12 16:52:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Collator fix

Browse Source
This commit is contained in:
Jake Poznanski 2025-06-29 19:52:53 +00:00
parent 12b5cc3101
commit dc7fff5bf7
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
olmocr/train/train.py
View File

@ -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'])
}