diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index d8f13560fe..46c9aac787 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -124,6 +124,7 @@ def _torchao_runtime_unavailable(): import sys try: from core._torchao_stub import is_win32_rocm, _STUB_SENTINEL + if is_win32_rocm(): return True _tao = sys.modules.get("torchao") diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 4785ed1a7d..fa57b08874 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -356,5 +356,9 @@ def test_installer_registers_finder_once(monkeypatch): assert len(finders) == 1 finally: sys.meta_path[:] = meta_before - for k in [k for k in sys.modules if (k == "torchao" or k.startswith("torchao.")) and k not in tao_before]: + for k in [ + k + for k in sys.modules + if (k == "torchao" or k.startswith("torchao.")) and k not in tao_before + ]: del sys.modules[k]