[Fix] Fix edge ID exclusion not working in EdgeDataLoader (dmlc#3412)

lizeyan · Oct 14, 2021 · 5d4f6bc · 5d4f6bc
1 parent 8798872
commit 5d4f6bc
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/dgl/dataloading/dataloader.py b/python/dgl/dataloading/dataloader.py
@@ -108,6 +108,7 @@ def __call__(self, frontier):
                         frontier, v, etype=k, store_ids=True)
                     new_eids[k] = F.gather_row(parent_eids[k], frontier.edges[k].data[EID])
             frontier.edata[EID] = new_eids
+        return frontier
 
 
 def _create_eid_excluder(exclude_eids, device):
@@ -388,7 +389,7 @@ def sample_blocks(self, g, seed_nodes, exclude_eids=None):
             if not self.exclude_edges_in_frontier:
                 eid_excluder = _create_eid_excluder(exclude_eids, self.output_device)
                 if eid_excluder is not None:
-                    eid_excluder(frontier)
+                    frontier = eid_excluder(frontier)
 
             block = transform.to_block(frontier, seed_nodes_out)
             if self.return_eids: