x

rapidsai · Dec 9, 2024 · a9d0c2d · a9d0c2d
1 parent 0635b77
commit a9d0c2d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage_wholegraph.py b/python/cugraph/cugraph/tests/data_store/test_gnn_feat_storage_wholegraph.py
@@ -76,6 +76,9 @@ def runtest(rank: int, world_size: int):
 @pytest.mark.skipif(
     get_cudart_version() < 11080, reason="not compatible with CUDA < 11.8"
 )
+@pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="PyTorch not installed with CUDA support"
+)
 def test_feature_storage_wholegraph_backend():
     world_size = torch.cuda.device_count()
     print("gpu count:", world_size)
@@ -94,6 +97,9 @@ def test_feature_storage_wholegraph_backend():
 @pytest.mark.skipif(
     get_cudart_version() < 11080, reason="not compatible with CUDA < 11.8"
 )
+@pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="PyTorch not installed with CUDA support"
+)
 def test_feature_storage_wholegraph_backend_mg():
     world_size = torch.cuda.device_count()
     print("gpu count:", world_size)