Delete device='cuda'

JerryL012 · Jun 26, 2023 · 3ee8ff7 · 3ee8ff7
1 parent 07957ba
commit 3ee8ff7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/cli_demo.py b/cli_demo.py
@@ -5,7 +5,7 @@
 import readline
 
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
-model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')
+model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
 model = model.eval()
 
 os_name = platform.system()

diff --git a/web_demo.py b/web_demo.py
@@ -3,7 +3,7 @@
 import mdtex2html
 
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
-model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')
+model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
 model = model.eval()
 
 """Override Chatbot.postprocess"""