quoccuongle/MultimodalOCR
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update interlm.py

Browse Source
This commit is contained in:
echo840
2024-02-28 17:18:21 +08:00
committed by GitHub
parent 39772991f3
commit 5611973faa
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
scripts/interlm.py
View File

@@ -49,7 +49,7 @@ def eval_worker(args, data, eval_id, output_queue):
torch.set_grad_enabled(False) torch.set_grad_enabled(False)
# init model and tokenizer # init model and tokenizer
model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True).cuda().eval() model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True,device_map=f'cuda:{eval_id}').eval()
tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
model.tokenizer = tokenizer model.tokenizer = tokenizer