From 43e4090ebe4d33fc5d441867c4055435d9b850a2 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 24 Jun 2026 01:09:40 +0000
Subject: [PATCH] Installer: add UNSLOTH_TORCH_INDEX_FAMILY override for torch
 wheel index

get_torch_index_url() (install.sh) and Get-TorchIndexUrl (install.ps1) pick
the PyTorch wheel index by probing the local GPU with nvidia-smi. That works
on a normal machine, but it has no good answer when there is no GPU to probe
at the time the index is chosen, for example inside a Docker image build or a
CI runner that builds wheels for a different target than it runs on. In that
situation the bash path falls back to cu126 (build hosts that expose
/proc/driver/nvidia but cannot run nvidia-smi) or cpu, baking the wrong
wheels into an image that is meant to run CUDA 12.8.

Add an optional UNSLOTH_TORCH_INDEX_FAMILY environment variable that names the
index path leaf directly (cu128, cu130, cu126, cpu, rocm6.4, ...). When set,
it short-circuits probing and pins the index deterministically. When unset or
empty the function behaves exactly as before, so native installs are
unaffected.
---
 install.ps1 |  5 +++++
 install.sh  | 11 +++++++++++
 2 files changed, 16 insertions(+)

diff --git a/install.ps1 b/install.ps1
index 745c7369ea..c35f7f447b 100644
--- a/install.ps1
+++ b/install.ps1
@@ -1939,6 +1939,11 @@ 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 override for hosts where GPU probing is impossible or must not
+        # happen (image builds, CI runners that build for a different target than
+        # they run on). Names the wheel index path leaf directly, e.g.
+        # UNSLOTH_TORCH_INDEX_FAMILY=cu128 | cu130 | cu126 | cpu | rocm6.4 | ...
+        if ($env:UNSLOTH_TORCH_INDEX_FAMILY) { return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY)" }
         if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
         try {
             $output = Invoke-NvidiaSmiBounded $NvidiaSmiExe
diff --git a/install.sh b/install.sh
index 5038de6264..823eee0de1 100755
--- a/install.sh
+++ b/install.sh
@@ -1962,6 +1962,17 @@ _has_usable_nvidia_gpu() {
 get_torch_index_url() {
     _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
     _base="${_base%/}"
+    # Explicit override for hosts where GPU probing is impossible or must not
+    # happen (Docker image builds, CI runners that build for a different
+    # target than they run on). Names the wheel index path leaf directly:
+    #   UNSLOTH_TORCH_INDEX_FAMILY=cu128 | cu130 | cu126 | cu124 | cu118 | cpu | rocm6.4 | ...
+    # Example: a Docker image that targets CUDA 12.8 is built on a host with no
+    # GPU and no nvidia-smi, so probing would otherwise fall back to cpu (or, on
+    # build hosts that leak /proc/driver/nvidia but cannot run nvidia-smi, cu126)
+    # and bake the wrong wheels. Setting this pins the index deterministically.
+    if [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then
+        echo "$_base/${UNSLOTH_TORCH_INDEX_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.