DUTwangzhijun commited on
Commit
653587a
·
1 Parent(s): fd3fd47

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +2 -2
model.py CHANGED
@@ -8,8 +8,8 @@ model_name = "DUTIRbionlp/Taiyi-LLM"
8
  model = AutoModelForCausalLM.from_pretrained(
9
  model_name,
10
  trust_remote_code=True,
11
- low_cpu_mem_usage=True,
12
- torch_dtype=torch.float16
13
  )
14
 
15
  model.eval()
 
8
  model = AutoModelForCausalLM.from_pretrained(
9
  model_name,
10
  trust_remote_code=True,
11
+ torch_dtype=torch.float16,
12
+ load_in_8bit = True
13
  )
14
 
15
  model.eval()