Merge pull request #541 from Anhforth/master

fixed bug
FlagAI-Open · Oct 7, 2023 · ae439ea · ae439ea
2 parents e1e0116 + df9c0c2
commit ae439ea
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/flagai/auto_model/auto_loader.py b/flagai/auto_model/auto_loader.py
@@ -266,8 +266,11 @@ def __init__(self,
                     bnb_4bit_compute_dtype=torch_dtype,
                 )
             if inference_mode:
-                model = AquilaForCausalLM.from_pretrained(download_path,low_cpu_mem_usage=low_cpu_mem_usage, torch_dtype=torch_dtype,
-                                                        quantization_config=quantization_config)
+                if qlora_dir:
+                    model = AquilaForCausalLM.from_pretrained(download_path,low_cpu_mem_usage=low_cpu_mem_usage, torch_dtype=torch_dtype,
+                                                            quantization_config=quantization_config)
+                else:
+                    model = AquilaForCausalLM.from_pretrained(download_path,low_cpu_mem_usage=low_cpu_mem_usage, torch_dtype=torch_dtype,)           
                 model.eval()
                 if not qlora_dir:
                     model.to(device)