From 7fe10c29efd6209bc63726b7dcb4c5293fe86720 Mon Sep 17 00:00:00 2001
From: Jeff Daily <jeff.daily@amd.com>
Date: Wed, 16 Jul 2025 21:17:35 +0000
Subject: [PATCH] adjust setup.py for ROCm

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

diff --git a/setup.py b/setup.py
index 0dde1e4..e562378 100644
--- a/setup.py
+++ b/setup.py
@@ -14,6 +14,7 @@
     CUDAExtension,
     BuildExtension,
     CUDA_HOME,
+    ROCM_HOME,
 )
 
 library_name = "extension_cpp"
@@ -30,7 +31,7 @@ def get_extensions():
     if debug_mode:
         print("Compiling in debug mode")
 
-    use_cuda = use_cuda and torch.cuda.is_available() and CUDA_HOME is not None
+    use_cuda = use_cuda and torch.cuda.is_available() and (CUDA_HOME is not None or ROCM_HOME is not None)
     extension = CUDAExtension if use_cuda else CppExtension
 
     extra_link_args = []