fix mse_v2 bug with PT1.13
@@ -18,6 +18,7 @@ import copy
import re
import numpy as np
from collections import UserDict
+from packaging.version import Version
from ...utils import logger
from ...utils.utility import LazyImport, CpuInfo
@@ -609,6 +610,9 @@ def get_mse_order_per_fp32(adaptor, model, example_inp, tune_cfg):
from torch.quantization.quantize_fx import prepare_fx,convert_fx
# do quantization
if adaptor.sub_module_list is None:
+ if adaptor.version.release >= Version("1.13.0").release: # pragma: no cover
+ tmp_model = prepare_fx(tmp_model, fx_op_cfgs, example_inp)
+ else:
tmp_model = prepare_fx(tmp_model, fx_op_cfgs,)
else:
PyTorch_FXAdaptor.prepare_sub_graph(adaptor.sub_module_list, fx_op_cfgs, \
@@ -658,6 +662,9 @@ def get_mse_order_per_fp32(adaptor, model, example_inp, tune_cfg):
from torch.quantization.quantize_fx import prepare_fx,convert_fx
# do quantization
if adaptor.sub_module_list is None:
+ if adaptor.version.release >= Version("1.13.0").release: # pragma: no cover
+ tmp_model = prepare_fx(tmp_model, fx_op_cfgs, example_inp)
+ else:
tmp_model = prepare_fx(tmp_model, fx_op_cfgs,)
else:
PyTorch_FXAdaptor.prepare_sub_graph(adaptor.sub_module_list, fx_op_cfgs, \
@@ -722,6 +729,9 @@ def get_mse_order_per_int8(adaptor, fp32_model, example_input, tune_cfg):
from torch.quantization.quantize_fx import prepare_fx,convert_fx
# do quantization
if adaptor.sub_module_list is None:
+ if adaptor.version.release >= Version("1.13.0").release: # pragma: no cover
+ tmp_model = prepare_fx(tmp_model, fx_op_cfgs, example_inp)
+ else:
tmp_model = prepare_fx(tmp_model, fx_op_cfgs,)
else:
PyTorch_FXAdaptor.prepare_sub_graph(adaptor.sub_module_list, fx_op_cfgs, \
