Delete device='cuda'

pull/23/head
duzx16 2023-06-26 13:32:14 +08:00
parent 07957bacf8
commit 3ee8ff7717
2 changed files with 2 additions and 2 deletions

View File

@ -5,7 +5,7 @@ from transformers import AutoTokenizer, AutoModel
import readline import readline
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) 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() model = model.eval()
os_name = platform.system() os_name = platform.system()

View File

@ -3,7 +3,7 @@ import gradio as gr
import mdtex2html import mdtex2html
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) 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() model = model.eval()
"""Override Chatbot.postprocess""" """Override Chatbot.postprocess"""