support new topk API in benchmark scripts

pyg-team · Jul 6, 2019 · e00adba · e00adba
1 parent 7198db0
commit e00adba
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/benchmark/kernel/sag_pool.py b/benchmark/kernel/sag_pool.py
@@ -39,7 +39,8 @@ def forward(self, data):
             xs += [global_mean_pool(x, batch)]
             if i % 2 == 0 and i < len(self.convs) - 1:
                 pool = self.pools[i // 2]
-                x, edge_index, _, batch, _ = pool(x, edge_index, batch=batch)
+                x, edge_index, _, batch, _, _ = pool(x, edge_index,
+                                                     batch=batch)
         x = self.jump(xs)
         x = F.relu(self.lin1(x))
         x = F.dropout(x, p=0.5, training=self.training)

diff --git a/benchmark/kernel/top_k.py b/benchmark/kernel/top_k.py
@@ -39,7 +39,8 @@ def forward(self, data):
             xs += [global_mean_pool(x, batch)]
             if i % 2 == 0 and i < len(self.convs) - 1:
                 pool = self.pools[i // 2]
-                x, edge_index, _, batch, _ = pool(x, edge_index, batch=batch)
+                x, edge_index, _, batch, _, _ = pool(x, edge_index,
+                                                     batch=batch)
         x = self.jump(xs)
         x = F.relu(self.lin1(x))
         x = F.dropout(x, p=0.5, training=self.training)