yujunjun/olmocr
1
0
Fork 0
mirror of https://github.com/allenai/olmocr.git synced 2025-10-13 17:22:13 +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']), 'input_ids': torch.stack(batch['input_ids']),
'attention_mask': torch.stack(batch['attention_mask']), 'attention_mask': torch.stack(batch['attention_mask']),
'labels': torch.stack(batch['labels']), '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']) 'image_grid_thw': torch.stack(batch['image_grid_thw'])
} }