From 50d867c13b22c22eda451ddb67bddb8159670f85 Mon Sep 17 00:00:00 2001 From: "Wang, Chang" Date: Thu, 4 Jul 2024 16:50:39 +0800 Subject: [PATCH] Update config import path for neural-compressor v2.6 (#801) Signed-off-by: changwangss --- optimum/intel/neural_compressor/trainer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/optimum/intel/neural_compressor/trainer.py b/optimum/intel/neural_compressor/trainer.py index ea9fe9a1cd..ff6dd95947 100644 --- 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