Fix edge_index construction in UPFD dataset (pyg-team#8413)

DamianSzwichtenberg · Nov 21, 2023 · 4f5100f · 4f5100f
1 parent af586eb
commit 4f5100f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -40,6 +40,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `edge_index` construction in the `UPFD` dataset ([#8413](https://github.com/pyg-team/pytorch_geometric/pull/8413))
 - Fixed TorchScript support in `AttentionalAggregation` and `DeepSetsAggregation` ([#8406](https://github.com/pyg-team/pytorch_geometric/pull/8406))
 - Fixed `GraphMaskExplainer` for GNNs with more than two layers ([#8401](https://github.com/pyg-team/pytorch_geometric/pull/8401))
 - Breaking Change: Properly initialize modules in `GATConv` depending on whether the input is bipartite or non-bipartite ([#8397](https://github.com/pyg-team/pytorch_geometric/pull/8397))

diff --git a/torch_geometric/datasets/upfd.py b/torch_geometric/datasets/upfd.py
@@ -144,7 +144,7 @@ def process(self):
         batch = torch.from_numpy(batch).to(torch.long)
 
         node_slice = cumsum(batch.bincount())
-        edge_slice = cumsum(batch[edge_index[0].bincount()])
+        edge_slice = cumsum(batch[edge_index[0]].bincount())
         graph_slice = torch.arange(y.size(0) + 1)
         self.slices = {
             'x': node_slice,