Skip to content

Commit 1d16103

Browse files
committed
Check neural-compressor version for export function
Signed-off-by: Cheng, Penghui <penghui.cheng@intel.com>
1 parent a141b3b commit 1d16103

File tree

2 files changed

+11
-2
lines changed

2 files changed

+11
-2
lines changed

optimum/intel/neural_compressor/quantization.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
from neural_compressor.model.onnx_model import ONNXModel
2929
from neural_compressor.model.torch_model import IPEXModel, PyTorchModel
3030
from neural_compressor.quantization import fit
31-
from neural_compressor.utils.export import torch_to_int8_onnx
3231
from packaging.version import parse
3332
from torch.utils.data import DataLoader, RandomSampler
3433
from transformers import (
@@ -80,6 +79,12 @@
8079
)
8180

8281

82+
if is_neural_compressor_version("<", "2.6"):
83+
from neural_compressor.experimental.export import torch_to_int8_onnx
84+
else:
85+
from neural_compressor.utils.export import torch_to_int8_onnx
86+
87+
8388
if is_itrex_available():
8489
if is_itrex_version("<", ITREX_MINIMUM_VERSION):
8590
raise ImportError(

optimum/intel/neural_compressor/trainer.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@
3636
from neural_compressor import training
3737
from neural_compressor.compression import DistillationCallbacks
3838
from neural_compressor.conf.pythonic_config import _BaseQuantizationConfig
39-
from neural_compressor.utils.export import torch_to_fp32_onnx, torch_to_int8_onnx
4039
from packaging import version
4140
from torch import nn
4241
from torch.utils.data import Dataset, RandomSampler
@@ -107,6 +106,11 @@
107106
if TYPE_CHECKING:
108107
from optimum.exporters.onnx import OnnxConfig
109108

109+
if is_neural_compressor_version("<", "2.6"):
110+
from neural_compressor.experimental.export import torch_to_fp32_onnx, torch_to_int8_onnx
111+
else:
112+
from neural_compressor.utils.export import torch_to_fp32_onnx, torch_to_int8_onnx
113+
110114

111115
__version__ = "4.22.2"
112116

0 commit comments

Comments
 (0)