pmelnechuk commited on
Commit
783a32a
verified
1 Parent(s): 2f94c75

Update src/model_load.py

Browse files
Files changed (1) hide show
  1. src/model_load.py +1 -1
src/model_load.py CHANGED
@@ -15,7 +15,7 @@ def load_model():
15
  model = AutoModelForCausalLM.from_pretrained(model_name,
16
  device_map="auto",
17
  quantization_config=BitsAndBytesConfig(load_in_8bit=True),
18
- torch_dtype=torch.float16,
19
  max_memory=max_memory)
20
 
21
  # Crear pipeline de generaci贸n de texto
 
15
  model = AutoModelForCausalLM.from_pretrained(model_name,
16
  device_map="auto",
17
  quantization_config=BitsAndBytesConfig(load_in_8bit=True),
18
+ torch_dtype="auto",
19
  max_memory=max_memory)
20
 
21
  # Crear pipeline de generaci贸n de texto