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

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

diff --git a/web_demo.py b/web_demo.py
index 1af24c9..63df6ba 100644
--- a/web_demo.py
+++ b/web_demo.py
@@ -4,7 +4,7 @@ import mdtex2html
 from utils import load_model_on_gpus
 
 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改为你实际的显卡数量
 # from utils import load_model_on_gpus
 # model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2)