From 39772991f3d1d3893c0f1ae110d5e714a08daeb6 Mon Sep 17 00:00:00 2001 From: echo840 <87795401+echo840@users.noreply.github.com> Date: Wed, 28 Feb 2024 17:17:33 +0800 Subject: [PATCH] Update interlm2.py --- scripts/interlm2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/interlm2.py b/scripts/interlm2.py index 7b920bd..3554b2b 100644 --- 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))):