From 8253e921cc237a458f9f61be78c35aa8ef4a9ea1 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 1 Jul 2025 00:18:22 -0700
Subject: [PATCH] Update setup.py

---
 setup.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/setup.py b/setup.py
index fc4c113166..b411d4d18c 100644
--- a/setup.py
+++ b/setup.py
@@ -30,9 +30,7 @@ IS_COLAB = "COLAB_" not in "".join(os.environ.keys())
 HAS_TORCH = False
 HAS_CUDA  = True
 HAS_HIP   = False
-CUDA_HOME = None
-ROCM_HOME = None
-torch = None
+torch     = None
 if importlib.util.find_spec("torch") is not None:
     try:
         import torch
@@ -44,9 +42,13 @@ if importlib.util.find_spec("torch") is not None:
     except Exception as e:
         print(f"Unsloth: Importing torch failed with error = {str(e)}.")
         HAS_TORCH = False
+        CUDA_HOME = None
+        ROCM_HOME = None
+else:
+    CUDA_HOME = None
+    ROCM_HOME = None
 pass
 
-
 def load_module_from_path(module_name, path):
     spec = importlib.util.spec_from_file_location(module_name, path)
     module = importlib.util.module_from_spec(spec)