clean-up

neuralmagic · Oct 13, 2024 · 34fecdb · 34fecdb
1 parent 4cf8508
commit 34fecdb
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 7 deletions.
diff --git a/src/compressed_tensors/quantization/lifecycle/apply.py b/src/compressed_tensors/quantization/lifecycle/apply.py
@@ -246,7 +246,6 @@ def apply_quantization_status(model: Module, status: QuantizationStatus):
     if current_status < status >= QuantizationStatus.CALIBRATION > current_status:
         # only quantize weights up front when our end goal state is calibration,
         # weight quantization parameters are already loaded for frozen/compressed
-        # TODO: to be removed from compressed-tensors
         quantize_weights_upfront = status == QuantizationStatus.CALIBRATION
         model.apply(
             lambda module: set_module_for_calibration(

diff --git a/src/compressed_tensors/quantization/lifecycle/frozen.py b/src/compressed_tensors/quantization/lifecycle/frozen.py
@@ -45,11 +45,7 @@ def freeze_module_quantization(module: Module):
         delattr(module, "input_observer")
     if hasattr(module, "weight_observer") and not scheme.weights.dynamic:
         delattr(module, "weight_observer")
-    if (
-        hasattr(module, "output_observer")
-        and not is_kv_cache_quant_scheme(scheme)
-        and not scheme.output_activations.dynamic
-    ):
+    if hasattr(module, "output_observer") and not scheme.output_activations.dynamic:
         delattr(module, "output_observer")
 
     module.quantization_status = QuantizationStatus.FROZEN
diff --git a/src/compressed_tensors/quantization/quant_args.py b/src/compressed_tensors/quantization/quant_args.py
@@ -121,7 +121,7 @@ def get_observer(self):
             self.observer = None
             return self.observer
 
-        return self.observer
+        return Observer.load_from_registry(self.observer, quantization_args=self)
 
     # TODO: update to be removed into llm-compressor
     def get_kv_cache(self):