fix test_serialization_big_model_inference failed

ZX-ModelCloud · Dec 21, 2024 · 9746143 · 9746143
1 parent 99b5f14
commit 9746143
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/tests/quantization/gptq/test_gptq.py b/tests/quantization/gptq/test_gptq.py
@@ -265,7 +265,16 @@ def test_serialization_big_model_inference(self):
         """
         with tempfile.TemporaryDirectory() as tmpdirname:
             self.quantized_model.save_pretrained(tmpdirname)
-            quantized_model_from_saved = AutoModelForCausalLM.from_pretrained(tmpdirname, device_map=self.device_map)
+
+            quantization_config = GPTQConfig(
+                bits=self.bits,
+                group_size=self.group_size,
+                desc_act=self.desc_act,
+                sym=self.sym,
+                use_exllama=self.use_exllama,
+            )
+            quantized_model_from_saved = AutoModelForCausalLM.from_pretrained(tmpdirname, device_map=self.device_map,
+                                                                              quantization_config=quantization_config)
             self.check_inference_correctness(quantized_model_from_saved)