From 5611973faa27a2afb2d8144acdaaf68cb4a68d3e Mon Sep 17 00:00:00 2001
From: echo840 <87795401+echo840@users.noreply.github.com>
Date: Wed, 28 Feb 2024 17:18:21 +0800
Subject: [PATCH] Update interlm.py

---
 scripts/interlm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/interlm.py b/scripts/interlm.py
index 85eed5e..2420aa1 100644
--- a/scripts/interlm.py
+++ b/scripts/interlm.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)
     model.tokenizer = tokenizer