From 46bc57ce49bc301f2431067a01baadf84bf6fc27 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 1 Jul 2025 00:34:27 -0700
Subject: [PATCH] Update setup.py

---
 setup.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/setup.py b/setup.py
index 192aea7b7f..e4e855912c 100644
--- a/setup.py
+++ b/setup.py
@@ -26,6 +26,9 @@ ROOT_DIR = Path(__file__).parent
 UNSLOTH_TARGET_DEVICE = os.environ.get("UNSLOTH_TARGET_DEVICE", "cuda")
 IS_COLAB = "COLAB_" not in "".join(os.environ.keys())
 
+import torch
+from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
+
 # Try importing torch
 HAS_TORCH = False
 HAS_CUDA  = True