From 732eab22c8417d2b2237d0ea60a9a91475620d92 Mon Sep 17 00:00:00 2001
From: zxgov <64576649+zx2021@users.noreply.github.com>
Date: Fri, 30 Jun 2023 10:25:01 +0800
Subject: [PATCH] Update api.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

使用更先进的方式加载模型
---
 api.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/api.py b/api.py
index e4236b7..b006614 100644
--- a/api.py
+++ b/api.py
@@ -51,7 +51,7 @@ async def create_item(request: Request):
 
 if __name__ == '__main__':
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
-    model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
+    model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')#.cuda()
     # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量
     # model_path = "THUDM/chatglm2-6b"
     # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)