From 71a9f3940aa07d2985d8a2ee9e1f914d0576f8ac Mon Sep 17 00:00:00 2001 From: "Wang, Mengni" Date: Tue, 9 Apr 2024 14:49:50 +0800 Subject: [PATCH] Fix ut failure for GPU backend (#1712) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- neural_compressor/adaptor/onnxrt.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/neural_compressor/adaptor/onnxrt.py b/neural_compressor/adaptor/onnxrt.py index a207fd78372..68da138f9fa 100644 --- a/neural_compressor/adaptor/onnxrt.py +++ b/neural_compressor/adaptor/onnxrt.py @@ -127,7 +127,9 @@ def __init__(self, framework_specific_info): self.work_space = framework_specific_info["workspace_path"] self.reduce_range = ( - framework_specific_info["reduce_range"] if "reduce_range" in framework_specific_info else not CpuInfo().vnni + framework_specific_info["reduce_range"] + if framework_specific_info.get("reduce_range", None) is not None + else not CpuInfo().vnni ) self.benchmark = GLOBAL_STATE.STATE == MODE.BENCHMARK os.makedirs(self.work_space, exist_ok=True)