yujunjun/olmocr
1
0
Fork 0
mirror of https://github.com/allenai/olmocr.git synced 2025-11-14 01:19:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix

Browse Source
This commit is contained in:
Jake Poznanski 2025-08-26 20:06:52 +00:00
parent 98f4d62d1e
commit 8327da2415
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
olmocr/train/grpo_train.py
View File

@ -473,6 +473,13 @@ def main():
trust_remote_code=True, trust_remote_code=True,
) )
# Load model
logger.info(f"Loading model: {args.model_name}")
if "Qwen2-VL" in args.model_name:
model_class = Qwen2VLForConditionalGeneration
else:
model_class = Qwen2_5_VLForConditionalGeneration
model = model_class.from_pretrained( model = model_class.from_pretrained(
args.model_name, args.model_name,
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,