pmelnechuk commited on
Commit
97dea46
·
verified ·
1 Parent(s): 896b328

Update src/model_load.py

Browse files
Files changed (1) hide show
  1. src/model_load.py +1 -1
src/model_load.py CHANGED
@@ -14,7 +14,7 @@ def load_model():
14
  tokenizer = AutoTokenizer.from_pretrained(model_name)
15
  model = AutoModelForCausalLM.from_pretrained(model_name,
16
  torch_dtype=torch.float16,
17
- device_map="auto"
18
  quantization_config=BitsAndBytesConfig(load_in_8bit=True),
19
  llm_int8_enable_fp32_cpu_offload=True)
20
 
 
14
  tokenizer = AutoTokenizer.from_pretrained(model_name)
15
  model = AutoModelForCausalLM.from_pretrained(model_name,
16
  torch_dtype=torch.float16,
17
+ device_map="auto",
18
  quantization_config=BitsAndBytesConfig(load_in_8bit=True),
19
  llm_int8_enable_fp32_cpu_offload=True)
20