Update tests

pyg-team · Dec 11, 2024 · 9c5c19e · 9c5c19e
1 parent 6197e8f
commit 9c5c19e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/test/nn/conv/test_gat_conv.py b/test/nn/conv/test_gat_conv.py
@@ -281,7 +281,8 @@ def test_remove_diag_sparse_tensor():
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
     edge_index2 = torch.tensor([[1, 2, 3], [0, 1, 1]])
 
-    adj1 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 4))
-    adj2 = SparseTensor.from_edge_index(edge_index2, sparse_sizes=(4, 4))
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj1 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 4))
+        adj2 = SparseTensor.from_edge_index(edge_index2, sparse_sizes=(4, 4))
 
-    assert torch_sparse.remove_diag(adj1.t()) == adj2.t()
+        assert torch_sparse.remove_diag(adj1.t()) == adj2.t()