Merge pull request #43 from teddykoker/cuda_fail_msg

Better error message if CUDA extension is not available
teddykoker · Dec 30, 2021 · 87f7715 · 87f7715
2 parents 6b81e6c + 8de04e1
commit 87f7715
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/setup.py b/setup.py
@@ -51,7 +51,7 @@ def ext_modules():
 
 setup(
     name="torchsort",
-    version="0.1.7",
+    version="0.1.8",
     description="Differentiable sorting and ranking in PyTorch",
     author="Teddy Koker",
     url="https://github.com/teddykoker/torchsort",

diff --git a/torchsort/ops.py b/torchsort/ops.py
@@ -26,10 +26,18 @@
     from .isotonic_cuda import isotonic_l2 as isotonic_l2_cuda
     from .isotonic_cuda import isotonic_l2_backward as isotonic_l2_backward_cuda
 except ImportError:
-    isotonic_l2_cuda = None
-    isotonic_kl_cuda = None
-    isotonic_l2_backward_cuda = None
-    isotonic_kl_backward_cuda = None
+
+    def _error(*args, **kwargs):
+        raise ImportError(
+            "You are trying to use the torchsort CUDA extension, but it looks like it is not available."
+            " Make sure you have the CUDA toolchain installed, and reinstall torchsort with `pip install --force-reinstall --no-cache-dir torchsort`"
+            " to rebuild the extension."
+        )
+
+    isotonic_l2_cuda = _error
+    isotonic_kl_cuda = _error
+    isotonic_l2_backward_cuda = _error
+    isotonic_kl_backward_cuda = _error
 
 
 def soft_rank(values, regularization="l2", regularization_strength=1.0):