fix: pass AMD torch install status via env var to suppress false warning

setup.ps1 now sets UNSLOTH_ROCM_TORCH_INSTALLED=1 after a successful AMD
wheel install. install_python_stack.py reads this at the top of
_ensure_rocm_torch() to skip both the subprocess probe and the warning --
no re-import of torch needed, and the warning message now correctly says
'could not be auto-installed' rather than 'must be installed manually'.
This commit is contained in:
LeoBorcherding 2026-05-06 20:37:44 -05:00
commit 7fbdce171c
2 changed files with 12 additions and 2 deletions

View file

@ -286,6 +286,12 @@ def _ensure_rocm_torch() -> None:
Uses pip_install() to respect uv, constraints, and --python targeting.
"""
global _rocm_windows_torch_installed
# setup.ps1 sets this env var when it successfully installs AMD wheels
# before calling install_python_stack.py, so we can skip the subprocess
# probe and avoid reinstalling what was just installed.
if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
_rocm_windows_torch_installed = True
return
if IS_MACOS:
return
@ -1191,11 +1197,11 @@ def install_python_stack() -> int:
if _win_amd_gpu and not _rocm_windows_torch_installed:
_safe_print(
_dim(" Note:"),
"AMD GPU detected on Windows. ROCm-enabled PyTorch must be",
"AMD GPU detected but ROCm PyTorch could not be auto-installed.",
)
_safe_print(
" " * 8,
"installed manually. See: https://docs.unsloth.ai/get-started/install-and-update/amd",
"Manual install may be required. See: https://docs.unsloth.ai/get-started/install-and-update/amd",
)
# 3. Extra dependencies