From 2c95b33ec2dd29c5eaf44ce3647498b9abdc2aee Mon Sep 17 00:00:00 2001 From: Akihiro Nitta Date: Thu, 27 Jul 2023 22:51:18 +0000 Subject: [PATCH] Update tests --- test/loader/test_link_neighbor_loader.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/loader/test_link_neighbor_loader.py b/test/loader/test_link_neighbor_loader.py index 4da61d966cac..fcdb0ee6f425 100644 --- a/test/loader/test_link_neighbor_loader.py +++ b/test/loader/test_link_neighbor_loader.py @@ -224,8 +224,9 @@ def test_temporal_homo_link_neighbor_loader(): ) for sample in loader: - assert sample.edge_label_index.size() == (batch_size,) + assert sample.edge_label_index.size() == (2, batch_size) assert sample.edge_label_time.size() == (batch_size,) + assert sample.edge_label.size() == (batch_size,) assert torch.all( sample.edge_time <= sample.edge_label_time ), "The target time should be later than all timestamps in the subgraph"