[Bug][Model] Bugfix in PinSAGE code (#1533)

0b902d03 · Quan (Andy) Gan · GitHub · 7b6af87a · 0b902d03
Unverified Commit 0b902d03 authored May 18, 2020 by Quan (Andy) Gan Committed by GitHub May 18, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

examples/pytorch/pinsage/sampler.py examples/pytorch/pinsage/sampler.py +3 -1

No files found.
--- a/examples/pytorch/pinsage/sampler.py
+++ b/examples/pytorch/pinsage/sampler.py
@@ -6,7 +6,9 @@ from torch.utils.data import IterableDataset, DataLoader
 def compact_and_copy(frontier, seeds):
    block = dgl.to_block(frontier, seeds)
    for col, data in frontier.edata.items():
-        block.edata[col] = data
+        if col == dgl.EID:
+            continue
+        block.edata[col] = data[block.edata[dgl.EID]]
    return block
 class ItemToItemBatchSampler(IterableDataset):