Skip to content

Commit 55b4cae

Browse files
authored
Fix torch check (#6402)
1 parent 0a4457c commit 55b4cae

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

op_builder/gds.py

+7-1
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,13 @@ def extra_ldflags(self):
3636
return super().extra_ldflags() + ['-lcufile']
3737

3838
def is_compatible(self, verbose=False):
39-
import torch.utils.cpp_extension
39+
try:
40+
import torch.utils.cpp_extension
41+
except ImportError:
42+
if verbose:
43+
self.warning("Please install torch if trying to pre-compile GDS")
44+
return False
45+
4046
CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME
4147
CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64")
4248
gds_compatible = self.has_function(funcname="cuFileDriverOpen",

0 commit comments

Comments
 (0)