From 10ba3dde84d00742f3635c48db09d6eccf0ec8bb Mon Sep 17 00:00:00 2001
From: Alex Morehead <alex.morehead@gmail.com>
Date: Wed, 4 Sep 2024 17:28:13 -0500
Subject: [PATCH] Handle an edge case where `CUDA_HOME` is not defined on ROCm
 systems (#6488)

* Handles an edge case when building `gds` where `CUDA_HOME` is not
defined on ROCm systems
---
 op_builder/gds.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/op_builder/gds.py b/op_builder/gds.py
index 727ebdf48372..d815ccea387a 100644
--- a/op_builder/gds.py
+++ b/op_builder/gds.py
@@ -49,6 +49,11 @@ def is_compatible(self, verbose=False):
             return False
 
         CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME
+        if CUDA_HOME is None:
+            if verbose:
+                self.warning("Please install torch CUDA if trying to pre-compile GDS with CUDA")
+            return False
+
         CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64")
         gds_compatible = self.has_function(funcname="cuFileDriverOpen",
                                            libraries=("cufile", ),