hahafofo commited on
Commit
f2f92a4
1 Parent(s): 1911897
Files changed (1) hide show
  1. utils/chatglm.py +6 -16
utils/chatglm.py CHANGED
@@ -81,22 +81,12 @@ class ChatGLM(BasePredictor):
81
  trust_remote_code=True,
82
  resume_download=True
83
  )
84
- if self.device == 'cuda':
85
- model = AutoModel.from_pretrained(
86
- model_name,
87
- trust_remote_code=True,
88
- resume_download=True
89
- ).half().to(self.device)
90
- else:
91
- model = AutoModel.from_pretrained(
92
- "THUDM/chatglm-6b-int8",
93
- trust_remote_code=True,
94
- resume_download=True,
95
- low_cpu_mem_usage=True,
96
- torch_dtype=torch.float16
97
- if self.device == 'cuda' else torch.float32,
98
- device_map={'': self.device}
99
- )
100
 
101
  model = model.eval()
102
  self.model = model
 
81
  trust_remote_code=True,
82
  resume_download=True
83
  )
84
+
85
+ model = AutoModel.from_pretrained(
86
+ model_name,
87
+ trust_remote_code=True,
88
+ resume_download=True
89
+ ).half().to(self.device)
 
 
 
 
 
 
 
 
 
 
90
 
91
  model = model.eval()
92
  self.model = model