diff --git a/cli_demo.py b/cli_demo.py
index 4aa35d3..586218e 100644
--- a/cli_demo.py
+++ b/cli_demo.py
@@ -5,7 +5,7 @@ from transformers import AutoTokenizer, AutoModel
 import readline
 
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
-model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
+model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')#.cuda()
 # 多显卡支持，使用下面两行代替上面一行，将num_gpus改为你实际的显卡数量
 # from utils import load_model_on_gpus
 # model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2)