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
microsoft · Sep 4, 2024 · 10ba3dd · 10ba3dd
1 parent 662a421
commit 10ba3dd
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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", ),