conditional import

ea7df4db · rusty1s · 66eef3b3 · ea7df4db
Commit ea7df4db authored Nov 15, 2018 by rusty1s
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

torch_cluster/sample.py torch_cluster/sample.py +2 -1

No files found.
--- a/torch_cluster/sample.py
+++ b/torch_cluster/sample.py
 import torch

-from sample_cuda import farthest_point_sampling, query_radius, query_knn
+if torch.cuda.is_available():
+    from sample_cuda import farthest_point_sampling, query_radius, query_knn


 def batch_slices(batch, sizes=False, include_ends=True):