From 77f61e677175225333237af28075b346c3dcde6e Mon Sep 17 00:00:00 2001 From: Logan Adams Date: Tue, 27 Aug 2024 13:02:07 -0700 Subject: [PATCH 1/2] Revert "Fix torch check (#6402)" This reverts commit 55b4cae80fc1479ac802b91b4b8955da444016ae. --- op_builder/gds.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/op_builder/gds.py b/op_builder/gds.py index 01c2d5a245d1..e024674e01d8 100644 --- a/op_builder/gds.py +++ b/op_builder/gds.py @@ -36,13 +36,7 @@ def extra_ldflags(self): return super().extra_ldflags() + ['-lcufile'] def is_compatible(self, verbose=False): - try: - import torch.utils.cpp_extension - except ImportError: - if verbose: - self.warning("Please install torch if trying to pre-compile GDS") - return False - + import torch.utils.cpp_extension CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64") gds_compatible = self.has_function(funcname="cuFileDriverOpen", From eb37cacf229ece71580a30bec723bf361f4ba82f Mon Sep 17 00:00:00 2001 From: Logan Adams Date: Tue, 27 Aug 2024 13:02:44 -0700 Subject: [PATCH 2/2] Revert "Revert "Fix torch check (#6402)"" This reverts commit 77f61e677175225333237af28075b346c3dcde6e. --- op_builder/gds.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/op_builder/gds.py b/op_builder/gds.py index e024674e01d8..01c2d5a245d1 100644 --- a/op_builder/gds.py +++ b/op_builder/gds.py @@ -36,7 +36,13 @@ def extra_ldflags(self): return super().extra_ldflags() + ['-lcufile'] def is_compatible(self, verbose=False): - import torch.utils.cpp_extension + try: + import torch.utils.cpp_extension + except ImportError: + if verbose: + self.warning("Please install torch if trying to pre-compile GDS") + return False + CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64") gds_compatible = self.has_function(funcname="cuFileDriverOpen",