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 14:03:37 +08:00
parent b35782c46e
commit 9d439304d4
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api.py
View File

@ -51,6 +51,6 @@ 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, 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)