From f542c231dc7089a5aedc67c10c600cd268981496 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 12 Apr 2026 22:09:27 +0000
Subject: [PATCH] Fix setup.ps1 version fast-path skipping ROCm repair on
 up-to-date installs

When an existing Windows AMD user runs `unsloth studio update` and the
unsloth package is already at the latest version, the version fast-path
set $SkipPythonDeps = $true, which bypassed the entire ROCm install
block. This left the user permanently stuck on CPU-only torch even
though the stale-venv check correctly identified the cpu->rocm mismatch.

Fix: track $_NeedRocmRepair flag from the stale-venv check and use it
to override $SkipPythonDeps so the ROCm wheel download still runs.

Also gate the stale-venv expectedTorchTag on $_NoTorch so --no-torch
users on AMD don't see a misleading "will repair to ROCm" message when
the actual $CuTag is "cpu".
---
 studio/setup.ps1 | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/studio/setup.ps1 b/studio/setup.ps1
index be8f4081e5..0828eb81e4 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -1510,10 +1510,15 @@ if (Test-Path $VenvDir -PathType Container) {
         $shouldRebuild = $true
     }
 
+    # $_NeedRocmRepair is used later to override the version fast-path
+    # ($SkipPythonDeps) so the ROCm install block still runs even when
+    # the unsloth package is already at the latest version.
+    $_NeedRocmRepair = $false
+
     if (-not $shouldRebuild) {
         if ($HasNvidiaSmi) {
             $expectedTorchTag = Get-PytorchCudaTag
-        } elseif ($HasAmdGpu) {
+        } elseif ($HasAmdGpu -and -not $_NoTorch) {
             $expectedTorchTag = "rocm"
         } else {
             $expectedTorchTag = "cpu"
@@ -1525,8 +1530,9 @@ if (Test-Path $VenvDir -PathType Container) {
             # setup.ps1 cannot recreate the venv on its own.  Instead, keep
             # the venv and let the ROCm install block below repair torch
             # in-place -- the end state is the same but no work is lost.
-            if ($HasAmdGpu -and $installedTorchTag -eq "cpu") {
+            if ($HasAmdGpu -and -not $_NoTorch -and $installedTorchTag -eq "cpu") {
                 substep "CPU-only torch detected on AMD host; will repair to ROCm in place..." "Yellow"
+                $_NeedRocmRepair = $true
             } else {
                 $shouldRebuild = $true
             }
@@ -1607,9 +1613,11 @@ if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") {
         $LatestVer = "$($pypiJson.info.version)".Trim()
     } catch { }
 
-    if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) {
+    if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer) -and -not $_NeedRocmRepair) {
         step "python" "$_PkgName $InstalledVer is up to date"
         $SkipPythonDeps = $true
+    } elseif ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer) -and $_NeedRocmRepair) {
+        substep "$_PkgName $InstalledVer is current; repairing torch to ROCm..."
     } elseif ($InstalledVer -and $LatestVer) {
         substep "$_PkgName $InstalledVer -> $LatestVer available, updating..."
     } elseif (-not $LatestVer) {