fix gptq

huggingface · Nov 17, 2023 · 6650bcb · 6650bcb
1 parent 0092945
commit 6650bcb
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/optimum_benchmark/backends/pytorch/backend.py b/optimum_benchmark/backends/pytorch/backend.py
@@ -242,8 +242,9 @@ def prepare_for_inference(self, input_shapes: Dict[str, int], **kwargs) -> None:
 
         if (self.config.quantization_scheme == "gptq" and self.config.quantization_config.get("desc_act", None)) or (
             hasattr(self.pretrained_config, "quantization_config")
-            and self.pretrained_config.quantization_config.get("desc_act", None)
-            and self.pretrained_config.quantization_config.get("quant_method", None) == "gptq"
+            and self.pretrained_config.quantization_config.quant_method == "gptq"
+            and hasattr(self.pretrained_config.quantization_config, "desc_act")
+            and self.pretrained_config.quantization_config.desc_act
         ):
             LOGGER.info("\t+ Setting GPTQ's max_input_length")
             from auto_gptq import exllama_set_max_input_length