Revert "Fix gptq exllamav2 check (#152)"

This reverts commit 5bf349d.
huggingface · Mar 13, 2024 · 6c89a82 · 6c89a82
1 parent 5bf349d
commit 6c89a82
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/optimum_benchmark/backends/pytorch/backend.py b/optimum_benchmark/backends/pytorch/backend.py
@@ -278,13 +278,15 @@ def is_awq_quantized(self) -> bool:
     def is_exllamav2(self) -> bool:
         return (self.is_gptq_quantized or self.is_awq_quantized) and (
             (
-                getattr(self.pretrained_config, "quantization_config", None) is not None
-                and getattr(self.pretrained_config.quantization_config, "exllama_config", None) is not None
-                and self.pretrained_config.quantization_config.exllama_config.get("exllama_version", None) == 2
+                hasattr(self.pretrained_config, "quantization_config")
+                and hasattr(self.pretrained_config.quantization_config, "exllama_config")
+                and "exllama_version" in self.pretrained_config.quantization_config.exllama_config
+                and self.pretrained_config.quantization_config.exllama_config["exllama_version"] == 2
             )
             or (
-                self.config.quantization_config.get("exllama_config", None) is not None
-                and self.config.quantization_config.exllama_config.get("exllama_version", None) == 2
+                hasattr(self.quantization_config, "exllama_config")
+                and "exllama_version" in self.quantization_config.exllama_config
+                and self.quantization_config.exllama_config["exllama_version"] == 2
             )
         )
 

diff --git a/optimum_benchmark/trackers/energy.py b/optimum_benchmark/trackers/energy.py
@@ -56,13 +56,13 @@ def __sub__(self, other: "Energy") -> "Energy":
         """Enables subtraction of two Energy instances using the '-' operator."""
         if self.unit != other.unit:
             raise ValueError("Energy units must match to perform subtraction")
-
+        
         return Energy(
             cpu=self.cpu - other.cpu,
             gpu=self.gpu - other.gpu,
             ram=self.ram - other.ram,
             total=self.total - other.total,
-            unit=self.unit,
+            unit=self.unit
         )