diff --git a/tests/neural_compressor/test_optimization.py b/tests/neural_compressor/test_optimization.py
index 15f450e19b..408a11f217 100644
--- a/tests/neural_compressor/test_optimization.py
+++ b/tests/neural_compressor/test_optimization.py
@@ -56,7 +56,7 @@
     INCSeq2SeqTrainer,
     INCStableDiffusionPipeline,
 )
-from optimum.intel.neural_compressor import WeightOnlyQuantConfig
+from intel_extension_for_transformers.transformers.utils.quantization_config import WeightOnlyQuantConfig
 from optimum.intel.utils.constant import DIFFUSION_WEIGHTS_NAME
 from optimum.onnxruntime import ORTModelForCausalLM, ORTModelForSequenceClassification
 from optimum.pipelines import ORT_SUPPORTED_TASKS