Update interlm2.py

2024-02-28 17:17:33 +08:00
parent e335694a91
commit 39772991f3
1 changed files with 1 additions and 1 deletions
--- a/scripts/interlm2.py
+++ b/scripts/interlm2.py
@@ -49,7 +49,7 @@ def eval_worker(args, data, eval_id, output_queue):
    torch.set_grad_enabled(False)

    # 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)

    for i in tqdm(range(len(data))):