Update quantization.py

intel · May 12, 2023 · a43bbd2 · a43bbd2
1 parent 06ec2d9
commit a43bbd2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/neural_compressor/quantization.py b/neural_compressor/quantization.py
@@ -160,10 +160,10 @@ def eval_func(model):
 
     if strategy_name == "mse_v2":
         if not (conf.framework.startswith("tensorflow")\
-                 or conf.framework == 'pytorch_fx'): # pragma: no cover
+                 or conf.framework in ['pytorch_fx', 'onnxruntime']): # pragma: no cover
             strategy_name = "basic"
             logger.warning(f"MSE_v2 does not support {conf.framework} now, use basic instead.")
-            logger.warning("Only tensorflow, pytorch_fx is supported by MSE_v2 currently.")
+            logger.warning("Only tensorflow, pytorch_fx, onnxruntime is supported by MSE_v2 currently.")
     assert strategy_name in STRATEGIES, "Tuning strategy {} is NOT supported".format(strategy_name)
 
     logger.info(f"Start {strategy_name} tuning.")