fix the test.

dmlc · Aug 21, 2024 · 57425e3 · 57425e3
1 parent e1150a1
commit 57425e3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/python/pytorch/graphbolt/impl/test_torch_based_feature_store.py b/tests/python/pytorch/graphbolt/impl/test_torch_based_feature_store.py
@@ -82,8 +82,8 @@ def test_torch_based_feature(in_memory):
         # Test get the size and count of the entire feature.
         assert feature_a.size() == torch.Size([3])
         assert feature_b.size() == torch.Size([2, 2])
-        assert feature_a.count() == a.size(0)
-        assert feature_b.count() == b.size(0)
+        assert feature_a.count() == 1
+        assert feature_b.count() == 3
 
         # Test get metadata of the feature.
         assert feature_a.metadata() == metadata