Update optimum/intel/neural_compressor/configuration.py

PenghuiCheng · echarlaix · web-flow · commit 5d90b5245650 · 2024-03-13T17:03:40.000+08:00
Co-authored-by: Ella Charlaix &lt;80481427+echarlaix@users.noreply.github.com&gt;
diff --git a/optimum/intel/neural_compressor/configuration.py b/optimum/intel/neural_compressor/configuration.py
@@ -35,7 +35,7 @@ class INCConfig(BaseConfig):
 
     def __init__(
         self,
-        quantization=None,
+        quantization: Optional[Union[Dict, _BaseQuantizationConfig, "WeightOnlyQuantConfig"]] = None,
         pruning: Optional[Union[Dict, _BaseQuantizationConfig]] = None,
         distillation: Optional[Union[Dict, _BaseQuantizationConfig]] = None,
         save_onnx_model: bool = False,
