Test labels in LinkNeighborLoader

test/loader/test_link_neighbor_loader.py

@@ -155,3 +155,34 @@ def test_heterogeneous_link_neighbor_loader_loop(directed):
         edge_label_index = batch['paper', 'paper'].edge_label_index
         edge_label_index = unique_edge_pairs(edge_label_index)
         assert len(edge_index | edge_label_index) == len(edge_index)
+
+
+def test_link_neighbour_loader_labels():
+
+    torch.manual_seed(12345)
+
+    edge_index = get_edge_index(100, 100, 500)
+    data = Data(edge_index=edge_index, x=torch.arange(100))
+
+    loader = LinkNeighborLoader(
+        data,
+        num_neighbors=[-1] * 2,
+        batch_size=10,
+        neg_sampling_ratio=1,
+    )
+
+    for batch in loader:
+        assert all(batch.edge_label[:10] == 1)
+        assert all(batch.edge_label[10:] == 0)
+
+    loader = LinkNeighborLoader(
+        data,
+        num_neighbors=[-1] * 2,
+        batch_size=10,
+        edge_label=torch.ones(500).type(torch.long),
+        neg_sampling_ratio=1,
+    )
+
+    for batch in loader:
+        assert all(batch.edge_label[:10] == 2)
+        assert all(batch.edge_label[10:] == 0)

torch_geometric/loader/link_neighbor_loader.py

@@ -170,7 +170,7 @@ class LinkNeighborLoader(torch.utils.data.DataLoader):
                  train_mask=[1368], val_mask=[1368], test_mask=[1368],
                  edge_label_index=[2, 128], edge_label=[128])
 
-    The rest of the functionality mirros that of
+    The rest of the functionality mirrors that of
     :class:`~torch_geometric.loader.NeighborLoader`, including support for
     heterogenous graphs.
 
@@ -213,7 +213,11 @@ class LinkNeighborLoader(torch.utils.data.DataLoader):
             :obj:`0` to :obj:`num_classes - 1`.
             After negative sampling, label :obj:`0` represents negative edges,
             and labels :obj:`1` to :obj:`num_classes` represent the labels of
-            positive edges. (default: :obj:`0.0`)
+            positive edges. (default: :obj:`0.0`).
+            Note that returned labels are :obj:`torch.float` when it is binary
+            classification (to facilitate ease-of-use with
+            :meth:`torch.binary_cross_entropy_loss`) and :obj:`torch.long`
+            when it is a multi-class classification problem.
         **kwargs (optional): Additional arguments of
             :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,
             :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.