Skip to content

Commit 68becc4

Browse files
Kaihui-intelpre-commit-ci[bot]CopilotXuehaoSun
authored
Reset accelerator when INC_TARGET_DEVICE is set in code (#2168)
Signed-off-by: Kaihui-intel <kaihui.tang@intel.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Sun, Xuehao <xuehao.sun@intel.com>
1 parent c9e265f commit 68becc4

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

neural_compressor/torch/utils/environ.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,10 @@ def get_accelerator(device_name="auto"):
150150
from neural_compressor.torch.utils.auto_accelerator import auto_detect_accelerator
151151

152152
accelerator = auto_detect_accelerator(device_name)
153+
inc_target_device = os.environ.get("INC_TARGET_DEVICE", None)
154+
if inc_target_device is not None and accelerator._name != inc_target_device.lower():
155+
auto_detect_accelerator.cache_clear()
156+
accelerator = auto_detect_accelerator(device_name)
153157
return accelerator
154158

155159

0 commit comments

Comments
 (0)