ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: setup.ps1 and install_python_stack.py now install ROCm torch on Windows

Browse source 
setup.ps1 was always setting CuTag='cpu' for non-NVIDIA hosts and installing
cpu-only PyTorch, overwriting the ROCm torch installed by install.ps1.
Adds the same AMD wheel selection logic (ROCm version detection, Python 3.12
check, 5-wheel install with --no-deps) to setup.ps1's torch install block.

install_python_stack.py: remove IS_WINDOWS guard from _ensure_rocm_torch()
call site so the Windows path in _ensure_rocm_torch() is reachable during
'unsloth studio update' as well.
This commit is contained in:
LeoBorcherding 2026-05-06 19:07:17 -05:00
commit b550948b39
2 changed files with 67 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

6
studio/install_python_stack.py
View file

@ -1140,12 +1140,12 @@ def install_python_stack() -> int:
# 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the
# venv received CPU-only torch (common when pip resolves torch from PyPI).
# Must come immediately after base packages so torch is present for inspection.
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
if not IS_MACOS and not NO_TORCH:
_progress("ROCm torch check")
_ensure_rocm_torch()
# Windows + AMD GPU: PyTorch does not publish ROCm wheels for Windows.
# Detect and warn so users know manual steps are needed for GPU training.
# Windows + AMD GPU: if ROCm torch was not installed (wrong Python version
# or unknown ROCm version), warn the user.
if IS_WINDOWS and not NO_TORCH and not _has_usable_nvidia_gpu():
# Validate actual AMD GPU presence (not just tool existence)
import re as _re_win