add test cases

test/nn/test_model.py

@@ -1,11 +1,9 @@
-import torch
 from relbench.base.task_base import TaskType
 from relbench.datasets.fake import FakeDataset
 from relbench.modeling.graph import (
     get_link_train_table_input,
     make_pkey_fkey_graph,
 )
-from relbench.modeling.loader import SparseTensor
 from relbench.modeling.utils import get_stype_proposal
 from relbench.tasks.amazon import UserItemPurchaseTask
 from torch_frame.config.text_embedder import TextEmbedderConfig
@@ -36,7 +34,7 @@ def test_idgnn(tmp_path):
     batch_size = 16
     train_loader = NeighborLoader(
         data,
-        num_neighbors=[128],
+        num_neighbors=[128, 128],
         time_attr="time",
         input_nodes=train_table_input.src_nodes,
         input_time=train_table_input.src_time,
@@ -46,7 +44,6 @@ def test_idgnn(tmp_path):
         shuffle=True,
     )
 
-    train_sparse_tensor = SparseTensor(train_table_input.dst_nodes[1])
     batch = next(iter(train_loader))
 
     assert len(batch[task.dst_entity_table].batch) > 0
@@ -55,22 +52,4 @@ def test_idgnn(tmp_path):
     model.train()
 
     out = model(batch, task.src_entity_table, task.dst_entity_table).flatten()
-
-    batch_size = batch[task.src_entity_table].batch_size
-
-    assert len(out) == batch_size
-
-    # Get ground-truth
-    input_id = batch[task.src_entity_table].input_id
-    # batch_offset, col_indices
-    src_batch, dst_index = train_sparse_tensor[input_id]
-
-    # Get target label
-    target = torch.isin(
-        batch[task.dst_entity_table].batch +
-        batch_size * batch[task.dst_entity_table].n_id,
-        src_batch + batch_size * dst_index,
-    ).float()
-    assert len(target) > 0
-    import pdb
-    pdb.set_trace()
+    assert len(out) == len(batch[task.dst_entity_table].n_id)