adapt test/data/test_batch.py

mova · Nov 24, 2023 · 8854100 · 8854100
1 parent 2474658
commit 8854100
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/test/data/test_batch.py b/test/data/test_batch.py
@@ -69,12 +69,12 @@ def test_batch_basic():
     assert str(batch[2]) == ("Data(x=[4], edge_index=[2, 6], y=[1], "
                              "string='3', array=[4], num_nodes=4)")
 
-    assert len(batch.index_select([1, 0])) == 2
-    assert len(batch.index_select(torch.tensor([1, 0]))) == 2
-    assert len(batch.index_select(torch.tensor([True, False]))) == 1
-    assert len(batch.index_select(np.array([1, 0], dtype=np.int64))) == 2
-    assert len(batch.index_select(np.array([True, False]))) == 1
-    assert len(batch[:2]) == 2
+    assert batch.index_select([1, 0]).num_graphs == 2
+    assert batch.index_select(torch.tensor([1, 0])).num_graphs == 2
+    assert batch.index_select(torch.tensor([True, False])).num_graphs == 1
+    assert batch.index_select(np.array([1, 0], dtype=np.int64)).num_graphs == 2
+    assert batch.index_select(np.array([True, False])).num_graphs == 1
+    assert batch[:2].num_graphs == 2
 
     data_list = batch.to_data_list()
     assert len(data_list) == 3