diff --git a/neural_compressor/utils/pytorch.py b/neural_compressor/utils/pytorch.py index ee49bfce855..4df4bd02162 100644 --- a/neural_compressor/utils/pytorch.py +++ b/neural_compressor/utils/pytorch.py @@ -210,7 +210,6 @@ def load_weight_only(checkpoint_dir, model, layer_wise=False): Returns: (object): quantized model """ - import neural_compressor # for eval(config['module_type']) from neural_compressor.adaptor.torch_utils.model_wrapper import MulLinear weights_file = os.path.join(os.path.abspath(os.path.expanduser(checkpoint_dir)), "best_model.pt") @@ -221,7 +220,7 @@ def load_weight_only(checkpoint_dir, model, layer_wise=False): for op_name, config in weight_only_config.items(): if config["dtype"] == "fp32": continue - if eval(config["module_type"]) == MulLinear: + if config["module_type"] == MulLinear.__module__ + "." + MulLinear.__name__: # op should be repleced by MulLinear module = util.fetch_module(model, op_name) new_module = MulLinear(module)