either load GPU or CPU libraries

torch_cluster/__init__.py

@@ -5,21 +5,24 @@ import torch
 
 __version__ = '1.5.9'
 
-suffix = 'cuda' if torch.cuda.is_available() else 'cpu'
-
 for library in [
         '_version', '_grid', '_graclus', '_fps', '_rw', '_sampler', '_nearest',
         '_knn', '_radius'
 ]:
-    torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
-        f'{library}_{suffix}', [osp.dirname(__file__)]).origin)
-
-if torch.cuda.is_available():  # pragma: no cover
-    cuda_version = torch.ops.torch_cluster.cuda_version()
-
-    if cuda_version == -1:
-        major = minor = 0
-    elif cuda_version < 10000:
+    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)
+    else:  # pragma: no cover
+        raise ImportError(f"Could not find module '{library}_cpu' in "
+                          f"{osp.dirname(__file__)}")
+
+cuda_version = torch.ops.torch_sparse.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])