Spaces:

pmelnechuk
/

MathQA

Sleeping

pmelnechuk commited on Feb 21

Commit

df93760

verified ·

1 Parent(s): 097094f

Update src/model_load.py

Files changed (1) hide show

src/model_load.py CHANGED Viewed

@@ -14,7 +14,7 @@ def load_model():
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name,
                                                  device_map="auto",
-                                                 quantization_config=BitsAndBytesConfig(load_in_4bit=True),
                                                  torch_dtype=torch.float16,
                                                  max_memory=max_memory)

     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name,
                                                  device_map="auto",
+                                                 quantization_config=BitsAndBytesConfig(load_in_8bit=True),
                                                  torch_dtype=torch.float16,
                                                  max_memory=max_memory)