Update config import path for neural-compressor v2.6 (#801)

Signed-off-by: changwangss <chang1.wang@intel.com>
huggingface · Jul 4, 2024 · 50d867c · 50d867c
1 parent 384dda5
commit 50d867c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optimum/intel/neural_compressor/trainer.py b/optimum/intel/neural_compressor/trainer.py
@@ -35,7 +35,6 @@
 import torch.distributed as dist
 from neural_compressor import training
 from neural_compressor.compression import DistillationCallbacks
-from neural_compressor.conf.pythonic_config import _BaseQuantizationConfig
 from packaging import version
 from torch import nn
 from torch.utils.data import Dataset, RandomSampler
@@ -107,8 +106,10 @@
     from optimum.exporters.onnx import OnnxConfig
 
 if is_neural_compressor_version("<", "2.6"):
+    from neural_compressor.conf.pythonic_config import _BaseQuantizationConfig
     from neural_compressor.experimental.export import torch_to_fp32_onnx, torch_to_int8_onnx
 else:
+    from neural_compressor.config import _BaseQuantizationConfig
     from neural_compressor.utils.export import torch_to_fp32_onnx, torch_to_int8_onnx