allow loading of CPU wheels in PyTorch CUDA installation

rusty1s · Aug 5, 2021 · 778f624 · 778f624
1 parent 66f87e5
commit 778f624
Showing 1 changed file with 38 additions and 39 deletions.
diff --git a/torch_scatter/__init__.py b/torch_scatter/__init__.py
@@ -6,50 +6,49 @@
 
 __version__ = '2.0.8'
 
-suffix = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-try:
-    for library in ['_version', '_scatter', '_segment_csr', '_segment_coo']:
-        torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
-            f'{library}_{suffix}', [osp.dirname(__file__)]).origin)
-except AttributeError as e:
-    if os.getenv('BUILD_DOCS', '0') != '1':
-        raise AttributeError(e)
-
-    from .placeholder import cuda_version_placeholder
-    torch.ops.torch_scatter.cuda_version = cuda_version_placeholder
-
-    from .placeholder import scatter_placeholder
-    torch.ops.torch_scatter.scatter_mul = scatter_placeholder
+for library in ['_version', '_scatter', '_segment_csr', '_segment_coo']:
+    cuda_spec = importlib.machinery.PathFinder().find_spec(
+        f'{library}_cuda', [osp.dirname(__file__)])
+    cpu_spec = importlib.machinery.PathFinder().find_spec(
+        f'{library}_cpu', [osp.dirname(__file__)])
+    spec = cuda_spec or cpu_spec
+    if spec is not None:
+        torch.ops.load_library(spec.origin)
+    elif os.getenv('BUILD_DOCS', '0') != '1':
+        raise ImportError(f"Could not find module '{library}_cpu' in "
+                          f"{osp.dirname(__file__)}")
+    else:
+        from .placeholder import cuda_version_placeholder
+        torch.ops.torch_scatter.cuda_version = cuda_version_placeholder
 
-    from .placeholder import scatter_arg_placeholder
-    torch.ops.torch_scatter.scatter_min = scatter_arg_placeholder
-    torch.ops.torch_scatter.scatter_max = scatter_arg_placeholder
+        from .placeholder import scatter_placeholder
+        torch.ops.torch_scatter.scatter_mul = scatter_placeholder
 
-    from .placeholder import segment_csr_placeholder
-    from .placeholder import segment_csr_arg_placeholder
-    from .placeholder import gather_csr_placeholder
-    torch.ops.torch_scatter.segment_sum_csr = segment_csr_placeholder
-    torch.ops.torch_scatter.segment_mean_csr = segment_csr_placeholder
-    torch.ops.torch_scatter.segment_min_csr = segment_csr_arg_placeholder
-    torch.ops.torch_scatter.segment_max_csr = segment_csr_arg_placeholder
-    torch.ops.torch_scatter.gather_csr = gather_csr_placeholder
+        from .placeholder import scatter_arg_placeholder
+        torch.ops.torch_scatter.scatter_min = scatter_arg_placeholder
+        torch.ops.torch_scatter.scatter_max = scatter_arg_placeholder
 
-    from .placeholder import segment_coo_placeholder
-    from .placeholder import segment_coo_arg_placeholder
-    from .placeholder import gather_coo_placeholder
-    torch.ops.torch_scatter.segment_sum_coo = segment_coo_placeholder
-    torch.ops.torch_scatter.segment_mean_coo = segment_coo_placeholder
-    torch.ops.torch_scatter.segment_min_coo = segment_coo_arg_placeholder
-    torch.ops.torch_scatter.segment_max_coo = segment_coo_arg_placeholder
-    torch.ops.torch_scatter.gather_coo = gather_coo_placeholder
+        from .placeholder import segment_csr_placeholder
+        from .placeholder import segment_csr_arg_placeholder
+        from .placeholder import gather_csr_placeholder
+        torch.ops.torch_scatter.segment_sum_csr = segment_csr_placeholder
+        torch.ops.torch_scatter.segment_mean_csr = segment_csr_placeholder
+        torch.ops.torch_scatter.segment_min_csr = segment_csr_arg_placeholder
+        torch.ops.torch_scatter.segment_max_csr = segment_csr_arg_placeholder
+        torch.ops.torch_scatter.gather_csr = gather_csr_placeholder
 
-if torch.cuda.is_available():  # pragma: no cover
-    cuda_version = torch.ops.torch_scatter.cuda_version()
+        from .placeholder import segment_coo_placeholder
+        from .placeholder import segment_coo_arg_placeholder
+        from .placeholder import gather_coo_placeholder
+        torch.ops.torch_scatter.segment_sum_coo = segment_coo_placeholder
+        torch.ops.torch_scatter.segment_mean_coo = segment_coo_placeholder
+        torch.ops.torch_scatter.segment_min_coo = segment_coo_arg_placeholder
+        torch.ops.torch_scatter.segment_max_coo = segment_coo_arg_placeholder
+        torch.ops.torch_scatter.gather_coo = gather_coo_placeholder
 
-    if cuda_version == -1:
-        major = minor = 0
-    elif cuda_version < 10000:
+cuda_version = torch.ops.torch_scatter.cuda_version()
+if torch.cuda.is_available() and cuda_version != -1:  # pragma: no cover
+    if cuda_version < 10000:
         major, minor = int(str(cuda_version)[0]), int(str(cuda_version)[2])
     else:
         major, minor = int(str(cuda_version)[0:2]), int(str(cuda_version)[3])