From 138c0949b4d94994a388627594371b055c12526f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 26 Jun 2026 04:59:32 +0000
Subject: [PATCH 1/9] install: let UNSLOTH_TORCH_INDEX_FAMILY / _URL override
 CUDA wheel detection

get_torch_index_url (and the studio-update mirror _detect_cuda_torch_index_url)
chose the torch wheel family solely by probing the host GPU, with no override.
In a headless / container / CI build the host driver is visible via the
/proc/driver/nvidia/gpus fallback but nvidia-smi cannot report a CUDA version,
so the function fell back to its cu126 default and installed the wrong wheels
(e.g. a cu128 image got cu126 torch).

Add an explicit override checked before any probing, in both the shell installer
and the Python studio-update path:
  - UNSLOTH_TORCH_INDEX_URL   full index URL, used verbatim (wins)
  - UNSLOTH_TORCH_INDEX_FAMILY family (cpu, cu128, rocm6.4, ...) appended to the
                               mirror base (UNSLOTH_PYTORCH_MIRROR still honoured)

This matches how the published GPU images select CUDA -- vLLM and SGLang take the
CUDA version from an explicit build ARG rather than detecting it, and the Unsloth
Docker base image already pins the cu128 index directly. Desktop installs are
unchanged: with no override set, detection runs exactly as before.

Adds test_get_torch_index_url.sh cases for the override (family, full URL,
precedence, mirror base, trailing-slash strip, empty-ignored).
---
 install.sh                           | 15 ++++++++++++
 studio/install_python_stack.py       | 17 +++++++++----
 tests/sh/test_get_torch_index_url.sh | 36 ++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/install.sh b/install.sh
index 548e6f702a..8e30264df0 100755
--- a/install.sh
+++ b/install.sh
@@ -1985,6 +1985,21 @@ _has_usable_nvidia_gpu() {
 get_torch_index_url() {
     _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
     _base="${_base%/}"
+    # Explicit override -- skip ALL GPU probing when the caller pins the wheel
+    # index. Headless / container / CI builds (and anyone cross-installing for a
+    # different target) must not let the build host's GPU -- or the lack of one --
+    # decide the wheel family. This is the same "tell the build, don't ask the
+    # hardware" approach the Docker base image and vLLM/SGLang's Dockerfiles take.
+    # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); UNSLOTH_TORCH_INDEX_FAMILY
+    # is the convenience form (cpu, cu124, cu126, cu128, cu130, rocm6.4, ...)
+    # appended to the mirror base so UNSLOTH_PYTORCH_MIRROR is still honoured.
+    if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ]; then
+        echo "${UNSLOTH_TORCH_INDEX_URL%/}"; return
+    fi
+    if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then
+        _family="${UNSLOTH_TORCH_INDEX_FAMILY#/}"
+        echo "$_base/${_family%/}"; return
+    fi
     # macOS: always CPU (no CUDA support)
     case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac
     # Try nvidia-smi -- require the binary to actually list a usable GPU.
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 9060b57542..222cfacb72 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -933,11 +933,20 @@ def _detect_cuda_torch_index_url() -> str:
 
     Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update`
     repairs to the same wheel family a fresh `curl | sh` install would pick.
-    Probes nvidia-smi (PATH, then /usr/bin/nvidia-smi) and parses both the
-    legacy "CUDA Version:" and the newer "CUDA UMD Version:" spellings.
-    Defaults to cu126 when nvidia-smi is missing or the version is unreadable
-    (e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback).
+    Honours the same explicit overrides first (UNSLOTH_TORCH_INDEX_URL /
+    UNSLOTH_TORCH_INDEX_FAMILY) so a headless / container / CI install never lets
+    the host GPU decide the wheel family. Otherwise probes nvidia-smi (PATH, then
+    /usr/bin/nvidia-smi) and parses both the legacy "CUDA Version:" and the newer
+    "CUDA UMD Version:" spellings. Defaults to cu126 when nvidia-smi is missing or
+    the version is unreadable (e.g. NVIDIA detected only via the
+    /proc/driver/nvidia/gpus fallback).
     """
+    _override_url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
+    if _override_url:
+        return _override_url.rstrip("/")
+    _override_family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
+    if _override_family:
+        return f"{_PYTORCH_WHL_BASE}/{_override_family.strip('/')}"
     exe = shutil.which("nvidia-smi")
     if not exe and os.path.isfile("/usr/bin/nvidia-smi"):
         exe = "/usr/bin/nvidia-smi"
diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh
index 6656142625..b9d0c5e5a1 100755
--- a/tests/sh/test_get_torch_index_url.sh
+++ b/tests/sh/test_get_torch_index_url.sh
@@ -379,6 +379,42 @@ _result=$(run_func "$_dir" " -1 ")
 assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
 rm -rf "$_dir"
 
+# --- explicit overrides (headless / container / CI; no GPU probing) ----------
+# 39) UNSLOTH_TORCH_INDEX_FAMILY pins the family with no GPU present -> that
+#     family (not the cpu fallback detection would pick).
+_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
+assert_eq "family override (no GPU) -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
+
+# 40) Family override wins over a real detection: a host whose nvidia-smi reports
+#     12.6 still gets cu128. This is the exact Docker-build case (the builder sees
+#     the host driver but must publish a cu128 image).
+_dir=$(make_mock_smi "12.6")
+_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "$_dir")
+assert_eq "family override beats detected 12.6 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
+rm -rf "$_dir"
+
+# 41) UNSLOTH_TORCH_INDEX_URL is used verbatim and wins over detection.
+_dir=$(make_mock_smi "12.6")
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu999" run_func "$_dir")
+assert_eq "url override beats detection -> verbatim" "https://mirror.example.com/whl/cu999" "$_result"
+rm -rf "$_dir"
+
+# 42) Family override is appended to UNSLOTH_PYTORCH_MIRROR (mirror still honoured).
+_result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
+assert_eq "mirror + family override -> mirror/cu128" "https://mirror.example.com/whl/cu128" "$_result"
+
+# 43) Trailing slash in UNSLOTH_TORCH_INDEX_URL is stripped.
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128/" run_func "none")
+assert_eq "url override trailing slash stripped" "https://mirror.example.com/whl/cu128" "$_result"
+
+# 44) URL override takes precedence over family override.
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu130" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
+assert_eq "url override beats family override -> url" "https://mirror.example.com/whl/cu130" "$_result"
+
+# 45) An empty override is ignored (falls through to normal detection).
+_result=$(UNSLOTH_TORCH_INDEX_FAMILY="" UNSLOTH_TORCH_INDEX_URL="" run_func "none")
+assert_eq "empty overrides ignored -> detected cpu" "https://download.pytorch.org/whl/cpu" "$_result"
+
 rm -f "$_FUNC_FILE"
 rm -rf "$_FAKE_SMI_DIR"
 rm -rf "$_TOOLS_DIR"

From b02a609af5ef4627bb46c7f6ffd88d03ea23914d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 26 Jun 2026 05:28:36 +0000
Subject: [PATCH 2/9] install: make the torch-index override authoritative
 across ROCm paths

Address review feedback on the override added in this PR so a pinned index is
honoured everywhere, not just in get_torch_index_url:

- Skip the WSL ROCm bootstrap (root privilege + large downloads, probes
  /dev/dxg) when UNSLOTH_TORCH_INDEX_URL / _FAMILY is set; it previously ran
  before the override was consulted.
- Skip the Radeon/Strix rerouting (which re-probes the GPU and overwrites the
  resolved URL with repo.radeon.com / repo.amd.com) when the index is pinned, so
  an explicit ROCm override (e.g. UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4) is kept.
- install_python_stack.py: derive _TORCH_BACKEND from the override when
  UNSLOTH_TORCH_BACKEND is unset (standalone studio update), so _ensure_rocm_torch
  / _ensure_cuda_torch repair to the requested family instead of re-detecting.
- Strip ALL leading/trailing slashes in the shell override to match the Python
  side (avoids 404s on strict pip proxies).

Adds test cases for double-slash and leading/trailing-slash overrides.
---
 install.sh                           | 29 ++++++++++++++++++++++++----
 studio/install_python_stack.py       | 17 ++++++++++++++++
 tests/sh/test_get_torch_index_url.sh |  8 ++++++++
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/install.sh b/install.sh
index 8e30264df0..e3fdf10cdd 100755
--- a/install.sh
+++ b/install.sh
@@ -1994,11 +1994,18 @@ get_torch_index_url() {
     # is the convenience form (cpu, cu124, cu126, cu128, cu130, rocm6.4, ...)
     # appended to the mirror base so UNSLOTH_PYTORCH_MIRROR is still honoured.
     if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ]; then
-        echo "${UNSLOTH_TORCH_INDEX_URL%/}"; return
+        # Strip ALL trailing slashes (match the Python side's .rstrip("/") and the
+        # Strix mirror handling below) -- a double/triple-slash URL 404s on strict
+        # pip proxies (artifactory, sonatype).
+        _url="${UNSLOTH_TORCH_INDEX_URL}"
+        while [ "${_url%/}" != "$_url" ]; do _url="${_url%/}"; done
+        echo "$_url"; return
     fi
     if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then
-        _family="${UNSLOTH_TORCH_INDEX_FAMILY#/}"
-        echo "$_base/${_family%/}"; return
+        _family="${UNSLOTH_TORCH_INDEX_FAMILY}"
+        while [ "${_family#/}" != "$_family" ]; do _family="${_family#/}"; done
+        while [ "${_family%/}" != "$_family" ]; do _family="${_family%/}"; done
+        echo "$_base/$_family"; return
     fi
     # macOS: always CPU (no CUDA support)
     case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac
@@ -2396,7 +2403,16 @@ _maybe_bootstrap_rocm_wsl() {
     [ -n "$_rw_tmp" ] && rm -f "$_rw_tmp"
     return 0
 }
-_maybe_bootstrap_rocm_wsl || true
+# When the caller pins the wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY),
+# honour it everywhere downstream: skip the WSL ROCm bootstrap (which can run
+# sudo + large downloads after probing /dev/dxg) and the Radeon/Strix rerouting
+# below (which would re-probe the GPU and overwrite the pinned URL). A headless /
+# container / CI build must get exactly the index it asked for.
+_torch_index_pinned=false
+if [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ] || [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then
+    _torch_index_pinned=true
+fi
+[ "$_torch_index_pinned" = true ] || _maybe_bootstrap_rocm_wsl || true
 
 TORCH_INDEX_URL=$(get_torch_index_url)
 
@@ -2424,7 +2440,11 @@ esac
 # Auto-detect GPU for AMD ROCm based
 # get_torch_index_url must have chosen */rocm*
 # (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon".
+# Skipped entirely when the index is pinned: an explicit override (even a ROCm
+# one like UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4) must not be rerouted to the
+# Radeon/Strix repos by GPU probing.
 _amd_gpu_radeon=false
+if [ "$_torch_index_pinned" = false ]; then
 case "$TORCH_INDEX_URL" in
     */rocm*)
         if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \
@@ -2505,6 +2525,7 @@ case "$TORCH_INDEX_URL" in
         fi
         ;;
 esac
+fi  # _torch_index_pinned guard (Radeon + Strix reroute)
 _TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL")
 if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then
     _TAURI_TORCH_INDEX_FAMILY="radeon"
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 222cfacb72..a2c1e44556 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -1429,6 +1429,23 @@ def _infer_no_torch() -> bool:
 # GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown
 # (standalone `unsloth studio update` runs, where we re-detect normally).
 _TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower()
+# When install.sh did not run (standalone `unsloth studio update`) but the caller
+# pinned the wheel index explicitly, derive the backend from that override so the
+# CUDA/ROCm repair helpers honour it instead of re-probing the GPU and possibly
+# reinstalling a different family. Classify on the final URL/family segment,
+# mirroring install.sh's UNSLOTH_TORCH_BACKEND case.
+if not _TORCH_BACKEND:
+    _idx_override = (
+        os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
+        or os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
+    )
+    _idx_leaf = _idx_override.rstrip("/").rsplit("/", 1)[-1].lower()
+    if _idx_leaf.startswith(("rocm", "gfx")):
+        _TORCH_BACKEND = "rocm"
+    elif _idx_leaf == "cpu":
+        _TORCH_BACKEND = "cpu"
+    elif _idx_leaf.startswith("cu"):
+        _TORCH_BACKEND = "cuda"
 
 
 def _torch_step_label(suffix: str) -> str:
diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh
index b9d0c5e5a1..0c538f0a12 100755
--- a/tests/sh/test_get_torch_index_url.sh
+++ b/tests/sh/test_get_torch_index_url.sh
@@ -415,6 +415,14 @@ assert_eq "url override beats family override -> url" "https://mirror.example.co
 _result=$(UNSLOTH_TORCH_INDEX_FAMILY="" UNSLOTH_TORCH_INDEX_URL="" run_func "none")
 assert_eq "empty overrides ignored -> detected cpu" "https://download.pytorch.org/whl/cpu" "$_result"
 
+# 46) ALL trailing slashes are stripped from a URL override (not just one).
+_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128///" run_func "none")
+assert_eq "url override double slash stripped" "https://mirror.example.com/whl/cu128" "$_result"
+
+# 47) Leading and trailing slashes stripped from a family override.
+_result=$(UNSLOTH_TORCH_INDEX_FAMILY="//cu128//" run_func "none")
+assert_eq "family override slashes stripped" "https://download.pytorch.org/whl/cu128" "$_result"
+
 rm -f "$_FUNC_FILE"
 rm -rf "$_FAKE_SMI_DIR"
 rm -rf "$_TOOLS_DIR"

From 997155d19367142ef989f764030a968849a542ac Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 26 Jun 2026 08:16:13 +0000
Subject: [PATCH 3/9] install: honor pinned torch index in CUDA/ROCm repair
 paths

Follow-up to the override work in this PR: the get_torch_index_url / install.sh
reroute already respect a pinned UNSLOTH_TORCH_INDEX_URL / _FAMILY, but the
Python repair helpers in install_python_stack.py still re-probed the GPU and
could overwrite the pinned family. Make the pin authoritative there too:

- _ensure_cuda_torch: an explicit cu* pin commits to CUDA wheels, so repair a
  ROCm-poisoned venv even when no NVIDIA GPU is visible here (headless /
  container / CI cross-install), instead of bailing on the GPU-presence gate.
- _ensure_rocm_torch: skip the AMD per-gfx (Strix) reroute when a ROCm index is
  pinned, and in the generic reinstall path install from the pinned URL verbatim
  rather than re-detecting the host ROCm version. gfx*/rocm7.2 indexes serve
  torch 2.11+, so select the 2.11 package specs for a gfx leaf.
- install.sh: raise the torch constraint to 2.11 for */gfx* indexes too, matching
  rocm7.2, so a pinned full-URL/family override that returns early keeps a valid
  constraint.

Add _explicit_torch_index_url / _explicit_rocm_torch_index_url helpers and tests
covering the no-GPU CUDA pin repair and the explicit gfx index honored verbatim.
---
 install.sh                                |  7 ++-
 studio/install_python_stack.py            | 73 ++++++++++++++++++-----
 tests/studio/install/test_cuda_repair.py  | 22 ++++++-
 tests/studio/install/test_rocm_support.py | 28 +++++++++
 4 files changed, 111 insertions(+), 19 deletions(-)

diff --git a/install.sh b/install.sh
index e3fdf10cdd..53b09cffab 100755
--- a/install.sh
+++ b/install.sh
@@ -2431,10 +2431,13 @@ case "$_torch_index_leaf" in
     *)          export UNSLOTH_TORCH_BACKEND="cuda" ;;
 esac
 
-# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it.
+# rocm7.2 and the AMD per-gfx indexes (repo.amd.com/.../gfxNNNN) ship torch
+# 2.11.0 -- adjust the constraint to allow it. This also covers a pinned full-URL
+# or family override (e.g. UNSLOTH_TORCH_INDEX_URL=.../gfx1151) that returns early
+# above and so never hits the Strix reroute that otherwise raises this constraint.
 # All other ROCm tags and CUDA stay within <2.11.0.
 case "$TORCH_INDEX_URL" in
-    */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
+    */rocm7.2|*/gfx*) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
 esac
 
 # Auto-detect GPU for AMD ROCm based
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index a2c1e44556..37c3c7d449 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -981,6 +981,30 @@ def _detect_cuda_torch_index_url() -> str:
     return f"{_PYTORCH_WHL_BASE}/{tag}"
 
 
+def _explicit_torch_index_url() -> "str | None":
+    """The wheel index URL pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY, else None.
+
+    Lets the CUDA/ROCm repair helpers honour the exact pinned family/URL instead
+    of re-probing the GPU. Mirrors install.sh::get_torch_index_url's override.
+    """
+    url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
+    if url:
+        return url.rstrip("/")
+    family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
+    if family:
+        return f"{_PYTORCH_WHL_BASE}/{family.strip('/')}"
+    return None
+
+
+def _explicit_rocm_torch_index_url() -> "str | None":
+    """The pinned wheel index URL when it names a ROCm family (rocm*/gfx*), else None."""
+    url = _explicit_torch_index_url()
+    if url is None:
+        return None
+    leaf = url.rstrip("/").rsplit("/", 1)[-1].lower()
+    return url if leaf.startswith(("rocm", "gfx")) else None
+
+
 def _ensure_cuda_torch() -> None:
     """Repair a venv whose torch is a ROCm build on an NVIDIA host.
 
@@ -1013,7 +1037,10 @@ def _ensure_cuda_torch() -> None:
         return
     # Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu()
     # covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent.
-    if not _has_usable_nvidia_gpu():
+    # An explicit CUDA wheel-index pin (headless / container / CI cross-install)
+    # commits to CUDA wheels regardless of whether a GPU is visible here, so it
+    # overrides the GPU-presence gate.
+    if not _has_usable_nvidia_gpu() and _explicit_torch_index_url() is None:
         return
 
     # Classify the installed torch: "hip" (ROCm build -- the poisoning
@@ -1256,7 +1283,9 @@ def _ensure_rocm_torch() -> None:
     # an incompatible wheel. Use HIP_VISIBLE_DEVICES for the runtime target.
     _strix_override_url: "str | None" = None
     _strix_override_pkgs: "tuple[str, str, str] | None" = None
-    if ver < (7, 2):
+    # An explicit ROCm wheel-index pin is authoritative: never auto-reroute it to
+    # the AMD per-gfx index (the caller already chose the family/URL).
+    if ver < (7, 2) and _explicit_rocm_torch_index_url() is None:
         gfx_codes = _detect_amd_gfx_codes()
         _strix_gfx = {"gfx1151", "gfx1150"}
         _detected_strix = _strix_gfx.intersection(gfx_codes)
@@ -1319,23 +1348,35 @@ def _ensure_rocm_torch() -> None:
         )
         rocm_torch_ready = True
     elif not has_hip_torch:
-        # Select best matching wheel tag (newest ROCm version <= installed)
-        tag = next(
-            (
-                t
-                for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
-                if ver >= (maj, mn)
-            ),
-            None,
-        )
+        # Honour an explicit ROCm wheel-index pin verbatim instead of re-detecting
+        # the host ROCm version; otherwise select the best wheel tag (newest ROCm
+        # version <= installed). gfx*/rocm7.2 indexes serve torch 2.11+, so match
+        # the constraints to the index leaf when overridden.
+        _override_idx = _explicit_rocm_torch_index_url()
+        if _override_idx is not None:
+            index_url = _override_idx
+            tag = index_url.rstrip("/").rsplit("/", 1)[-1].lower()
+        else:
+            tag = next(
+                (
+                    t
+                    for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True)
+                    if ver >= (maj, mn)
+                ),
+                None,
+            )
         if tag is None:
             print(f"   No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " f"skipping torch reinstall")
         else:
-            index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
-            print(f"   ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}")
-            _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get(
-                tag, _ROCM_TORCH_PKG_SPECS["_default"]
-            )
+            if _override_idx is None:
+                index_url = f"{_PYTORCH_WHL_BASE}/{tag}"
+            print(f"   ROCm torch -- installing from {index_url}")
+            if tag.startswith("gfx"):
+                _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["rocm7.2"]
+            else:
+                _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get(
+                    tag, _ROCM_TORCH_PKG_SPECS["_default"]
+                )
             pip_install(
                 f"ROCm torch ({tag})",
                 "--force-reinstall",
diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py
index cea4383268..38e9764ea4 100644
--- a/tests/studio/install/test_cuda_repair.py
+++ b/tests/studio/install/test_cuda_repair.py
@@ -64,15 +64,19 @@ def _run_cuda_repair(
     rocm_marker = False,
     smi_path = "/usr/bin/nvidia-smi",
     cvd = None,
+    index_family = None,
 ):
     """Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock.
 
-    cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it."""
+    cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it.
+    index_family sets UNSLOTH_TORCH_INDEX_FAMILY (the explicit wheel-index pin)."""
     env = {}
     if rocm_marker:
         env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1"
     if cvd is not None:
         env["CUDA_VISIBLE_DEVICES"] = cvd
+    if index_family is not None:
+        env["UNSLOTH_TORCH_INDEX_FAMILY"] = index_family
 
     def _which(name, *a, **k):
         if name == "nvidia-smi":
@@ -99,6 +103,9 @@ def _which(name, *a, **k):
             stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None)
         if cvd is None:
             stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        if index_family is None:
+            stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+            stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
         _ensure_cuda_torch()
     return mock_pip
 
@@ -128,6 +135,19 @@ def test_rocm_in_version_string_triggers_repair(self):
         mock_pip = _run_cuda_repair(torch_state = "hip")
         assert mock_pip.call_count == 1
 
+    def test_no_gpu_but_explicit_cuda_pin_repairs(self):
+        # Headless / container / CI cross-install: an explicit cu* index pin
+        # commits to CUDA wheels even though no NVIDIA GPU is visible here, so a
+        # ROCm-poisoned venv is still repaired (to the pinned family).
+        mock_pip = _run_cuda_repair(
+            nvidia = False,
+            backend = "cuda",
+            index_family = "cu128",
+            torch_state = "hip",
+        )
+        assert mock_pip.call_count == 1
+        assert "cu128" in _index_url(mock_pip)
+
 
 # No-op cases.
 
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index 4c50fa1c8e..b46967dffe 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -680,6 +680,34 @@ def test_rocm_72_selects_72_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip)
         torch_call = mock_pip.call_args_list[0]
         assert "rocm7.2" in str(torch_call)
 
+    @patch.object(stack_mod, "IS_WINDOWS", False)
+    @patch.object(stack_mod, "pip_install_try", return_value = True)
+    @patch.object(stack_mod, "pip_install")
+    @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
+    @patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
+    @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4))
+    def test_explicit_gfx_index_honored_and_skips_strix_reroute(
+        self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
+    ):
+        """An explicit gfx wheel-index pin is authoritative: install from it verbatim
+        with torch 2.11, and never re-probe gfx codes to second-guess it (host ROCm 6.4
+        would otherwise pick the rocm6.4 wheel / trigger the Strix re-route)."""
+        mock_probe = MagicMock()
+        mock_probe.returncode = 0
+        mock_probe.stdout = b"\n"  # cpu torch -> reinstall
+        env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"}
+        with patch.dict(stack_mod.os.environ, env, clear = False):
+            stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
+            with patch("os.path.isdir", return_value = True):
+                with patch("subprocess.run", return_value = mock_probe):
+                    # Would raise if the Strix block ran (it is skipped on an explicit pin).
+                    with patch.object(stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError):
+                        _ensure_rocm_torch()
+        assert mock_pip.call_count == 1
+        torch_call = str(mock_pip.call_args_list[0])
+        assert "gfx1151" in torch_call
+        assert "torch>=2.11.0,<2.12.0" in torch_call
+
     @patch.object(stack_mod, "IS_WINDOWS", False)
     @patch.object(stack_mod, "pip_install_try", return_value = True)
     @patch.object(stack_mod, "pip_install")

From 0c0e2cbbc71fc675713565ba57c71affc1a78ae4 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 26 Jun 2026 08:17:31 +0000
Subject: [PATCH 4/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/studio/install/test_rocm_support.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index b46967dffe..a900cffb90 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -701,7 +701,9 @@ def test_explicit_gfx_index_honored_and_skips_strix_reroute(
             with patch("os.path.isdir", return_value = True):
                 with patch("subprocess.run", return_value = mock_probe):
                     # Would raise if the Strix block ran (it is skipped on an explicit pin).
-                    with patch.object(stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError):
+                    with patch.object(
+                        stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
+                    ):
                         _ensure_rocm_torch()
         assert mock_pip.call_count == 1
         torch_call = str(mock_pip.call_args_list[0])

From 5736d5da9604838f476fa4a5eed67536ec941590 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 26 Jun 2026 08:33:48 +0000
Subject: [PATCH 5/9] install: honor torch-index override on the Windows
 installers too

The pinned-index work landed for install.sh and install_python_stack.py, but the
Windows installers still picked the wheel index from GPU probing. Extend the same
UNSLOTH_TORCH_INDEX_URL / _FAMILY contract so a pinned index wins on every platform:

- install.ps1: Get-TorchIndexUrl returns the pinned URL/family before nvidia-smi
  probing; the AMD ROCm reroute is skipped when the index is pinned, so an explicit
  cpu/cu* pin on an AMD host is not overwritten.
- studio/setup.ps1: add shared Get-PinnedTorchIndexUrl / Get-TorchIndexLeaf helpers;
  the stale-venv check, the install selection and the AMD reroute all honor the pin,
  and the CPU/CUDA install pulls from the resolved index URL.
- tests: parity test that all four installers read both override vars and the two
  Windows installers gate the AMD reroute on the pinned flag.
---
 install.ps1                                | 18 +++++++-
 studio/setup.ps1                           | 49 ++++++++++++++++++----
 tests/python/test_cross_platform_parity.py | 30 +++++++++++++
 3 files changed, 89 insertions(+), 8 deletions(-)

diff --git a/install.ps1 b/install.ps1
index f7f9540970..2a929c81d5 100644
--- a/install.ps1
+++ b/install.ps1
@@ -1939,6 +1939,17 @@ exit 0
     # Mirrors Get-PytorchCudaTag in setup.ps1.
     function Get-TorchIndexUrl {
         $baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
+        # Explicit pin -- skip ALL GPU probing when the caller names the wheel index
+        # (headless / CI / cross-install). Matches install.sh::get_torch_index_url and
+        # install_python_stack.py: UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim);
+        # UNSLOTH_TORCH_INDEX_FAMILY is the convenience leaf (cpu, cu128, rocm6.4, ...)
+        # appended to the mirror base so UNSLOTH_PYTORCH_MIRROR is still honoured.
+        if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
+            return $env:UNSLOTH_TORCH_INDEX_URL.Trim().TrimEnd('/')
+        }
+        if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
+            return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
+        }
         if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
         try {
             $output = Invoke-NvidiaSmiBounded $NvidiaSmiExe
@@ -2016,6 +2027,11 @@ exit 0
         } catch { return $null }
     }
 
+    # An explicit UNSLOTH_TORCH_INDEX_URL / _FAMILY pin is authoritative: the AMD
+    # ROCm reroute below must not rewrite it (e.g. a deliberate cpu pin on an AMD
+    # host, or a pinned ROCm family we already resolved in Get-TorchIndexUrl).
+    $TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or `
+                        (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY))
     $TorchIndexUrl = Get-TorchIndexUrl
 
     # ── GPU arch → newest compatible Windows ROCm wheel release ──
@@ -2027,7 +2043,7 @@ exit 0
     # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
     $ROCmIndexUrl = $null
     $ROCmTorchFloor = $null
-    if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
+    if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
         $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
         $archFamilyMap = @{
             "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all"  # RDNA 4
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index 38398eb5f5..907367a69a 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -402,6 +402,29 @@ function Get-PytorchCudaTag {
     return "cu126"
 }
 
+# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY), shared by the
+# stale-venv check and the install selection below so a pinned wheel index wins
+# over GPU probing -- matching install.sh, install.ps1 and install_python_stack.py.
+# UNSLOTH_TORCH_INDEX_URL is verbatim (full URL); _FAMILY is the leaf (cpu, cu128,
+# rocm6.4, ...) joined to the mirror base so UNSLOTH_PYTORCH_MIRROR is honoured.
+function Get-PinnedTorchIndexUrl {
+    if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
+        return $env:UNSLOTH_TORCH_INDEX_URL.Trim().TrimEnd('/')
+    }
+    if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
+        $base = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
+        return "$base/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
+    }
+    return $null
+}
+
+# The last path segment of a wheel index URL (cu128 / cpu / rocm6.4 / gfx1151).
+function Get-TorchIndexLeaf {
+    param([string]$Url)
+    if ([string]::IsNullOrWhiteSpace($Url)) { return $null }
+    return ($Url.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
+}
+
 # VS generator -> MSBuild BuildCustomizations dir; toolset tracks the VS major
 # (18->v180, 17->v170), defaulting to v170 when unparseable.
 function Get-VcBuildCustomizationsDir {
@@ -2536,7 +2559,8 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
     }
 
     if (-not $shouldRebuild) {
-        $expectedTorchTag = if ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" }
+        $_pinnedIdx = Get-PinnedTorchIndexUrl
+        $expectedTorchTag = if ($_pinnedIdx) { Get-TorchIndexLeaf $_pinnedIdx } elseif ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" }
         if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
             $shouldRebuild = $true
         }
@@ -2712,7 +2736,13 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir
 [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User')
 substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)"
 
-if ($HasNvidiaSmi) {
+# Explicit pin (URL or family) wins over GPU probing and suppresses the AMD
+# reroute below; matches install.sh / install.ps1 / install_python_stack.py.
+$PinnedTorchIndexUrl = Get-PinnedTorchIndexUrl
+$TorchIndexPinned = [bool]$PinnedTorchIndexUrl
+if ($PinnedTorchIndexUrl) {
+    $CuTag = Get-TorchIndexLeaf $PinnedTorchIndexUrl
+} elseif ($HasNvidiaSmi) {
     $CuTag = Get-PytorchCudaTag
 } else {
     $CuTag = "cpu"
@@ -2733,7 +2763,7 @@ $ROCmIndexUrl = $null
 # SDK -- which flips Studio out of chat-only (CHAT_ONLY) and enables Train/Export.
 # Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed
 # ROCm install still falls back to CPU below, so this is safe.
-if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
+if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
     $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
     $archFamilyMap = @{
         "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all"  # RDNA 4
@@ -2785,6 +2815,11 @@ if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
 
 $PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
 
+# A full UNSLOTH_TORCH_INDEX_URL pin is used verbatim; a family pin already set
+# $CuTag, so $PyTorchWhlBase/$CuTag is the requested family index. The CPU/CUDA
+# install branches below pull from this instead of re-joining mirror + tag.
+$TorchInstallIndexUrl = if ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" }
+
 $ROCmCpuFallback = $false
 if ($ROCmIndexUrl) {
     substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..."
@@ -2821,11 +2856,11 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
     $cpuForce = @()
     if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") }
     if ($script:UnslothVerbose) {
-        Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu"
+        Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl
         $torchInstallExit = $LASTEXITCODE
         $output = ""
     } else {
-        $output = Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" | Out-String
+        $output = Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl | Out-String
         $torchInstallExit = $LASTEXITCODE
     }
     if ($torchInstallExit -ne 0) {
@@ -2837,11 +2872,11 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
     substep "installing PyTorch with CUDA support ($CuTag)..."
     substep "(This download is ~2.8 GB -- may take a few minutes)"
     if ($script:UnslothVerbose) {
-        Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag"
+        Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl
         $torchInstallExit = $LASTEXITCODE
         $output = ""
     } else {
-        $output = Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" | Out-String
+        $output = Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl | Out-String
         $torchInstallExit = $LASTEXITCODE
     }
     if ($torchInstallExit -ne 0) {
diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py
index 666ea7ce10..0fae0c3255 100644
--- a/tests/python/test_cross_platform_parity.py
+++ b/tests/python/test_cross_platform_parity.py
@@ -10,6 +10,8 @@
 REPO_ROOT = Path(__file__).resolve().parents[2]
 INSTALL_SH = REPO_ROOT / "install.sh"
 INSTALL_PS1 = REPO_ROOT / "install.ps1"
+SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1"
+STACK_PY = REPO_ROOT / "studio" / "install_python_stack.py"
 
 
 class TestNoTorchBackendAutoInInstallSh:
@@ -180,3 +182,31 @@ def test_install_ps1_preserves_timeout_override(self):
         assert (
             '$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text
         ), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT"
+
+
+class TestTorchIndexOverrideParity:
+    """Every installer must honor UNSLOTH_TORCH_INDEX_URL / _FAMILY so a pinned wheel
+    index wins over GPU probing on all platforms (no asymmetric, per-OS coverage)."""
+
+    @pytest.mark.parametrize(
+        "path",
+        [INSTALL_SH, INSTALL_PS1, SETUP_PS1, STACK_PY],
+        ids = ["install.sh", "install.ps1", "setup.ps1", "install_python_stack.py"],
+    )
+    def test_installer_reads_override_env(self, path):
+        text = path.read_text(encoding = "utf-8")
+        for var in ("UNSLOTH_TORCH_INDEX_URL", "UNSLOTH_TORCH_INDEX_FAMILY"):
+            assert var in text, f"{path.name} does not honor {var}"
+
+    @pytest.mark.parametrize(
+        "path",
+        [INSTALL_PS1, SETUP_PS1],
+        ids = ["install.ps1", "setup.ps1"],
+    )
+    def test_amd_reroute_guarded_when_pinned(self, path):
+        # The AMD ROCm reroute must be skipped when the index is explicitly pinned,
+        # so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten.
+        text = path.read_text(encoding = "utf-8")
+        assert "TorchIndexPinned" in text, (
+            f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
+        )

From 5a017ecfea1d2de579dc78a50debe8140e2bd29f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 26 Jun 2026 08:34:29 +0000
Subject: [PATCH 6/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/python/test_cross_platform_parity.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py
index 0fae0c3255..ae1e247634 100644
--- a/tests/python/test_cross_platform_parity.py
+++ b/tests/python/test_cross_platform_parity.py
@@ -207,6 +207,6 @@ def test_amd_reroute_guarded_when_pinned(self, path):
         # The AMD ROCm reroute must be skipped when the index is explicitly pinned,
         # so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten.
         text = path.read_text(encoding = "utf-8")
-        assert "TorchIndexPinned" in text, (
-            f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
-        )
+        assert (
+            "TorchIndexPinned" in text
+        ), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"

From d2d5f90175960c00e2e2759c92ce166004ac127f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 26 Jun 2026 09:16:19 +0000
Subject: [PATCH 7/9] install: complete pinned-index handling for ROCm/Windows
 edge cases

Follow-ups to the override work flagged in review:

- install.ps1: a pinned gfx*/rocm>=7.2 index previously skipped the AMD reroute
  that sets the torch>=2.11 floor, so the generic install used torch>=2.4,<2.11
  and could resolve the known-bad _grouped_mm wheel. Route a pinned ROCm index
  through the ROCm install path with the 2.11 floor + companions, and guard the
  companion-spec lookup so a skipped reroute block cannot null-deref.
- studio/setup.ps1: the stale-venv check compared the installed flavor (cuXXX/cpu,
  with +rocm misread as cpu) against the raw pinned leaf (gfx1151 / rocm6.4), so a
  correct pinned ROCm venv was always marked stale. Classify +rocm wheels as the
  generic 'rocm' flavor and normalize a pinned rocm*/gfx* leaf to 'rocm' before
  comparing (cu* stays specific so cu126-vs-cu128 still rebuilds).
- install_python_stack.py: _ensure_cuda_torch now also reinstalls from a pinned
  CUDA index when the venv carries a CPU wheel (headless CPU-venv-to-CUDA
  cross-install via 'studio update'), not only when it finds a ROCm build.
- tests: parity assertions already cover all four installers honoring the override.
---
 install.ps1                    | 34 ++++++++++++++++++++++++++++++----
 studio/install_python_stack.py | 25 ++++++++++++++++++++-----
 studio/setup.ps1               | 17 ++++++++++++++++-
 3 files changed, 66 insertions(+), 10 deletions(-)

diff --git a/install.ps1 b/install.ps1
index 2a929c81d5..7e4e8b4aac 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2043,6 +2043,8 @@ exit 0
     # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
     $ROCmIndexUrl = $null
     $ROCmTorchFloor = $null
+    $PinnedRocmVisionSpec = $null
+    $PinnedRocmAudioSpec = $null
     if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
         $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
         $archFamilyMap = @{
@@ -2093,6 +2095,30 @@ exit 0
         }
     }
 
+    # An explicit gfx*/rocm pin skips the auto-reroute above, but the generic
+    # CPU/CUDA install below would use torch>=2.4,<2.11 and pull a known-bad wheel
+    # on the gfx115x/gfx120x/rocm>=7.2 indexes (the torch._C._grouped_mm null-ptr
+    # bug). Route a pinned ROCm index through the ROCm install path with the same
+    # 2.11 floor/companions the unpinned reroute derives from the gfx arch.
+    if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) {
+        $_pinLeaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLower()
+        $_pinRocm211 = $false
+        if ($_pinLeaf -match '^rocm(\d+)\.(\d+)') {
+            $_pinRocm211 = ([int]$Matches[1] -gt 7) -or ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -ge 2)
+        }
+        if ($_pinLeaf -like 'gfx*' -or $_pinRocm211) {
+            $ROCmIndexUrl = $TorchIndexUrl
+            $ROCmTorchFloor = "torch>=2.11.0,<2.12.0"
+            $PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0"
+            $PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0"
+            substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan"
+        } elseif ($_pinLeaf -like 'rocm*') {
+            # Older rocm (<=7.1) ships torch <2.11; route via the ROCm path with the
+            # default floor so the pinned family resolves its own wheels.
+            $ROCmIndexUrl = $TorchIndexUrl
+        }
+    }
+
     if ($ROCmIndexUrl) {
         $TorchIndexFamily = "rocm"
     } else {
@@ -2205,8 +2231,8 @@ exit 0
             $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
             # Pin the companions to match $torchSpec; bare names can resolve an
             # ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
-            $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
-            $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
+            $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
+            $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
             $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
             if ($torchInstallExit -ne 0) {
                 # Transient AMD-index failure: fall back to a CPU base so the install
@@ -2326,8 +2352,8 @@ exit 0
                     $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
                     # Pin companions like the fresh ROCm path (bare names can pull an
                     # ABI-incompatible torchvision/torchaudio from the per-arch index).
-                    $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
-                    $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
+                    $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
+                    $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
                     substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
                     $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
                     if ($torchFixExit -ne 0) {
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 37c3c7d449..ceb73951d8 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -1073,16 +1073,31 @@ def _ensure_cuda_torch() -> None:
     _marker_lines = [
         line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip()
     ]
-    if not _marker_lines or _marker_lines[-1] != "hip":
+    if not _marker_lines:
+        return
+    _marker = _marker_lines[-1]
+    # Reinstall CUDA torch when the venv carries a ROCm build on an NVIDIA host
+    # (the poisoning signature), OR when an explicit CUDA index is pinned but the
+    # venv still has a CPU wheel. The latter is the headless CPU-venv-to-CUDA
+    # cross-install (`studio update` with UNSLOTH_TORCH_INDEX_FAMILY=cu128): the
+    # update path preserves torch rather than preinstalling it from install.sh, so
+    # without this the explicit CUDA pin stays ineffective. A healthy CUDA torch,
+    # or a CPU wheel with no CUDA pin, is deliberate and left alone.
+    _pin = _explicit_torch_index_url()
+    _pinned_cuda = bool(_pin) and _pin.rstrip("/").rsplit("/", 1)[-1].lower().startswith("cu")
+    if _marker == "hip":
+        _why = "torch is a ROCm build on an NVIDIA host"
+    elif _marker == "cpu" and _pinned_cuda:
+        _why = "torch is a CPU build but an explicit CUDA index is pinned"
+    else:
         return  # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is
 
     index_url = _detect_cuda_torch_index_url()
     _torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
     print(
-        f"   torch is a ROCm build on an NVIDIA host -- reinstalling "
-        f"CUDA torch from {index_url}\n"
-        f"   (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch "
-        f"on a mixed AMD+NVIDIA host)"
+        f"   {_why} -- reinstalling CUDA torch from {index_url}\n"
+        f"   (set UNSLOTH_TORCH_BACKEND=rocm or cpu to keep a deliberate "
+        f"non-CUDA torch)"
     )
     pip_install(
         "CUDA torch repair",
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index 907367a69a..4b22579c3b 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -2540,6 +2540,11 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
             if ($finished -and $proc.ExitCode -eq 0 -and $torchVer) {
                 if ($torchVer -match '\+(cu\d+)') {
                     $installedTorchTag = $Matches[1]
+                } elseif ($torchVer -match '\+rocm') {
+                    # Any +rocm / gfx wheel -> generic "rocm" flavor. The exact ROCm
+                    # version is repaired later by install_python_stack.py; here we
+                    # only need the flavor so a correct ROCm venv is not marked stale.
+                    $installedTorchTag = "rocm"
                 } elseif ($torchVer -match '\+cpu') {
                     $installedTorchTag = "cpu"
                 } else {
@@ -2560,7 +2565,17 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
 
     if (-not $shouldRebuild) {
         $_pinnedIdx = Get-PinnedTorchIndexUrl
-        $expectedTorchTag = if ($_pinnedIdx) { Get-TorchIndexLeaf $_pinnedIdx } elseif ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" }
+        if ($_pinnedIdx) {
+            $_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx
+            # Normalize a pinned rocm*/gfx* leaf to the generic "rocm" flavor so it
+            # compares against the installed +rocm wheel (also "rocm"); cu*/cpu
+            # leaves stay specific so a cu126-vs-cu128 mismatch still rebuilds.
+            $expectedTorchTag = if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { "rocm" } else { $_pinLeaf }
+        } elseif ($HasNvidiaSmi) {
+            $expectedTorchTag = Get-PytorchCudaTag
+        } else {
+            $expectedTorchTag = "cpu"
+        }
         if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
             $shouldRebuild = $true
         }

From 4a5baba8cea526f367393f42b1e2b18d5285da69 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 26 Jun 2026 10:36:59 +0000
Subject: [PATCH 8/9] install: finish pinned ROCm/CUDA edge cases on Windows +
 repair path

Follow-ups to the previous round:

- studio/setup.ps1: a pinned gfx*/rocm>=7.2 index now routes through the ROCm
  install path with the 2.11 floor + companions (it previously fell through to the
  CUDA branch with bare torch/torchvision/torchaudio against the ROCm index). The
  CPU/CUDA fallback index is forced to the CPU wheel index when a ROCm index is
  active, so a failed pinned-ROCm install does not retry the ROCm mirror.
- studio/setup.ps1: the stale-venv check no longer treats an unrecognized pinned
  URL leaf (e.g. a PEP 503 mirror ending in /simple) as a torch flavor tag, which
  was marking a correct venv stale; cu*/cpu/rocm/gfx leaves are still compared.
- install.ps1: the post-failure CPU fallback uses an explicit CPU index instead of
  , which for a pinned ROCm index was the ROCm mirror itself (so the
  'fallback' just retried the failing index and aborted the installer).
- install_python_stack.py: _ensure_cuda_torch now also reinstalls when the venv's
  CUDA family differs from a pinned one (installed cu126 vs pinned cu128), not only
  CPU->CUDA; the probe reports the installed cuXXX tag for the comparison.
---
 install.ps1                    |  6 ++++-
 studio/install_python_stack.py | 28 +++++++++++++---------
 studio/setup.ps1               | 44 +++++++++++++++++++++++++++++++---
 3 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/install.ps1 b/install.ps1
index 7e4e8b4aac..edec79b29d 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2242,7 +2242,11 @@ exit 0
                 # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
                 # torch>= range, so without it uv would keep the ROCm build and only swap
                 # the companions -- a mismatched venv the flavor-repair block won't fix.
-                $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
+                # Use an explicit CPU index: in the unpinned AMD path $TorchIndexUrl is
+                # already */cpu, but for a pinned ROCm index it IS the ROCm mirror, so
+                # reusing it here would just retry the failing ROCm index, not fall back.
+                $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" }
+                $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $CpuFallbackIndexUrl }
                 if ($torchInstallExit -ne 0) {
                     Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
                     return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index ceb73951d8..0aa6116a08 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -1053,11 +1053,13 @@ def _ensure_cuda_torch() -> None:
                 sys.executable,
                 "-c",
                 (
-                    "import torch; "
+                    "import torch, re; "
                     "hip = getattr(torch.version, 'hip', '') or ''; "
                     "cuda = getattr(torch.version, 'cuda', '') or ''; "
                     "ver = getattr(torch, '__version__', '').lower(); "
-                    "print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))"
+                    "m = re.search(r'\\+(cu\\d+)', ver); "
+                    "marker = 'hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'); "
+                    "print(marker + '|' + (m.group(1) if m else ''))"
                 ),
             ],
             stdout = subprocess.PIPE,
@@ -1075,22 +1077,26 @@ def _ensure_cuda_torch() -> None:
     ]
     if not _marker_lines:
         return
-    _marker = _marker_lines[-1]
+    _marker, _, _installed_cu = _marker_lines[-1].partition("|")
     # Reinstall CUDA torch when the venv carries a ROCm build on an NVIDIA host
-    # (the poisoning signature), OR when an explicit CUDA index is pinned but the
-    # venv still has a CPU wheel. The latter is the headless CPU-venv-to-CUDA
-    # cross-install (`studio update` with UNSLOTH_TORCH_INDEX_FAMILY=cu128): the
-    # update path preserves torch rather than preinstalling it from install.sh, so
-    # without this the explicit CUDA pin stays ineffective. A healthy CUDA torch,
-    # or a CPU wheel with no CUDA pin, is deliberate and left alone.
+    # (the poisoning signature), or when an explicit CUDA index is pinned but the
+    # venv has the wrong family -- a CPU wheel, or a different cuXXX than pinned.
+    # This covers the headless cross-install (`studio update` with
+    # UNSLOTH_TORCH_INDEX_FAMILY=cu128): the update path preserves torch rather
+    # than preinstalling it from install.sh, so without this an explicit CUDA pin
+    # stays ineffective. A healthy CUDA torch matching the pin, or a CPU wheel
+    # with no CUDA pin, is deliberate and left alone.
     _pin = _explicit_torch_index_url()
-    _pinned_cuda = bool(_pin) and _pin.rstrip("/").rsplit("/", 1)[-1].lower().startswith("cu")
+    _pin_leaf = _pin.rstrip("/").rsplit("/", 1)[-1].lower() if _pin else ""
+    _pinned_cuda = _pin_leaf.startswith("cu")
     if _marker == "hip":
         _why = "torch is a ROCm build on an NVIDIA host"
     elif _marker == "cpu" and _pinned_cuda:
         _why = "torch is a CPU build but an explicit CUDA index is pinned"
+    elif _marker == "cuda" and _pinned_cuda and _installed_cu and _installed_cu != _pin_leaf:
+        _why = f"torch is {_installed_cu} but the pinned CUDA index is {_pin_leaf}"
     else:
-        return  # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is
+        return  # healthy CUDA torch matching the pin, or a deliberate CPU wheel
 
     index_url = _detect_cuda_torch_index_url()
     _torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index 4b22579c3b..ffbef0bac8 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -2565,18 +2565,29 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
 
     if (-not $shouldRebuild) {
         $_pinnedIdx = Get-PinnedTorchIndexUrl
+        $_expectedKnown = $true
         if ($_pinnedIdx) {
             $_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx
             # Normalize a pinned rocm*/gfx* leaf to the generic "rocm" flavor so it
             # compares against the installed +rocm wheel (also "rocm"); cu*/cpu
             # leaves stay specific so a cu126-vs-cu128 mismatch still rebuilds.
-            $expectedTorchTag = if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { "rocm" } else { $_pinLeaf }
+            if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') {
+                $expectedTorchTag = "rocm"
+            } elseif ($_pinLeaf -like 'cu*' -or $_pinLeaf -eq 'cpu') {
+                $expectedTorchTag = $_pinLeaf
+            } else {
+                # Custom index whose final segment is not a torch flavor (e.g. a
+                # PEP 503 mirror ending in /simple). We cannot infer the flavor, so
+                # trust the pinned URL and do not rebuild on a bogus tag comparison.
+                $_expectedKnown = $false
+                $expectedTorchTag = $installedTorchTag
+            }
         } elseif ($HasNvidiaSmi) {
             $expectedTorchTag = Get-PytorchCudaTag
         } else {
             $expectedTorchTag = "cpu"
         }
-        if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
+        if ($_expectedKnown -and $installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
             $shouldRebuild = $true
         }
     }
@@ -2828,12 +2839,39 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu
     }
 }
 
+# A pinned gfx*/rocm index skips the auto-reroute above; route it through the
+# ROCm install path with the same floor/companions the unpinned AMD path uses
+# (mirrors install.ps1). Otherwise the CUDA branch below installs bare torch /
+# torchvision / torchaudio from the ROCm index and resolves a known-bad <2.11
+# wheel or ABI-mismatched companions for gfx115x / gfx120x / rocm>=7.2.
+if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) {
+    $_pinLeaf = Get-TorchIndexLeaf $PinnedTorchIndexUrl
+    $_pinRocm211 = $false
+    if ($_pinLeaf -match '^rocm(\d+)\.(\d+)') {
+        $_pinRocm211 = ([int]$Matches[1] -gt 7) -or ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -ge 2)
+    }
+    if ($_pinLeaf -like 'gfx*' -or $_pinRocm211) {
+        $ROCmIndexUrl   = $PinnedTorchIndexUrl
+        $ROCmTorchSpec  = "torch>=2.11.0,<2.12.0"
+        $ROCmVisionSpec = "torchvision>=0.26.0,<0.27.0"
+        $ROCmAudioSpec  = "torchaudio>=2.11.0,<2.12.0"
+        substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchSpec" "Cyan"
+    } elseif ($_pinLeaf -like 'rocm*') {
+        $ROCmIndexUrl   = $PinnedTorchIndexUrl
+        $ROCmTorchSpec  = "torch"
+        $ROCmVisionSpec = "torchvision"
+        $ROCmAudioSpec  = "torchaudio"
+    }
+}
+
 $PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
 
 # A full UNSLOTH_TORCH_INDEX_URL pin is used verbatim; a family pin already set
 # $CuTag, so $PyTorchWhlBase/$CuTag is the requested family index. The CPU/CUDA
 # install branches below pull from this instead of re-joining mirror + tag.
-$TorchInstallIndexUrl = if ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" }
+# A pinned ROCm install goes through $ROCmIndexUrl; if that fails, the fallback
+# must use the CPU wheel index, not retry the ROCm mirror left in the pin.
+$TorchInstallIndexUrl = if ($ROCmIndexUrl) { "$PyTorchWhlBase/cpu" } elseif ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" }
 
 $ROCmCpuFallback = $false
 if ($ROCmIndexUrl) {

From f6c5e46dadae896f7c80e9a35707620ce8520b36 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 26 Jun 2026 11:36:33 +0000
Subject: [PATCH 9/9] install: keep the ROCm to CPU fallback install inside the
 retry-helper window

The pinned-ROCm CPU fallback computes an explicit CPU index, but the comment
explaining why it cannot reuse $TorchIndexUrl pushed the actual
Invoke-InstallCommandRetry / --force-reinstall call more than 600 chars past the
"ROCm PyTorch install failed" message, so test_pr5940_followups's window check
no longer saw the retry helper. Move the CPU-index computation and its comment
above the failure substep so the retrying force-reinstall stays adjacent to the
message. No behavior change: same explicit CPU index, same retry, same
--force-reinstall.
---
 install.ps1 | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/install.ps1 b/install.ps1
index edec79b29d..b895ec24d9 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2237,15 +2237,15 @@ exit 0
             if ($torchInstallExit -ne 0) {
                 # Transient AMD-index failure: fall back to a CPU base so the install
                 # still completes; Studio setup retries ROCm afterwards.
+                # Use an explicit CPU index: in the unpinned AMD path $TorchIndexUrl is
+                # already */cpu, but for a pinned ROCm index it IS the ROCm mirror, so
+                # reusing it here would just retry the failing ROCm index, not fall back.
+                $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" }
                 substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Studio setup retries ROCm." "Yellow"
                 # --force-reinstall: a failed ROCm install can leave an unpinned ROCm
                 # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
                 # torch>= range, so without it uv would keep the ROCm build and only swap
                 # the companions -- a mismatched venv the flavor-repair block won't fix.
-                # Use an explicit CPU index: in the unpinned AMD path $TorchIndexUrl is
-                # already */cpu, but for a pinned ROCm index it IS the ROCm mirror, so
-                # reusing it here would just retry the failing ROCm index, not fall back.
-                $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" }
                 $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $CpuFallbackIndexUrl }
                 if ($torchInstallExit -ne 0) {
                     Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red