Gaze/ChatGLM2-6B
1
1
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Remove device="cuda"

Browse Source
This commit is contained in:
duzx16 2023-06-26 22:22:20 +08:00
parent d09927fa52
commit 3444756a17
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
openai_api.py
View File

@ -157,7 +157,7 @@ async def predict(query: str, history: List[List[str]], model_id: str):
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, device="cuda")
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
model.eval()
uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)