Merge pull request #17 from myml/memory

fix: 修复chatglm模型被复制，显存占用过多
2023-04-07 09:43:49 +08:00 · 2023-04-07 09:43:49 +08:00 · 63d900607f
parent 51c44e3e0a bed03a6ff1
commit 63d900607f
1 changed files with 2 additions and 0 deletions
--- a/chatglm_llm.py
+++ b/chatglm_llm.py
@ -15,6 +15,8 @@ model = (
    .cuda()
 )

+tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
+model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()

 class ChatGLM(LLM):
    max_token: int = 10000