richardr1126 commited on
Commit
89b2762
·
1 Parent(s): 4cf036c

Back to a10g

Browse files
Files changed (1) hide show
  1. app.py +1 -2
app.py CHANGED
@@ -20,9 +20,8 @@ print(f"Starting to load the model {model_name}")
20
 
21
  m = AutoModelForCausalLM.from_pretrained(
22
  model_name,
23
- device_map='auto',
24
  #load_in_8bit=True,
25
- llm_int8_enable_fp32_cpu_offload=True
26
  )
27
 
28
  m.config.pad_token_id = m.config.eos_token_id
 
20
 
21
  m = AutoModelForCausalLM.from_pretrained(
22
  model_name,
23
+ device_map=0,
24
  #load_in_8bit=True,
 
25
  )
26
 
27
  m.config.pad_token_id = m.config.eos_token_id