Skip to content
Permalink
Browse files

support new topk API in benchmark scripts

  • Loading branch information...
rusty1s committed Jul 6, 2019
1 parent 7198db0 commit e00adba648a62d11ab066b02615a4479dd7b503f
Showing with 4 additions and 2 deletions.
  1. +2 −1 benchmark/kernel/sag_pool.py
  2. +2 −1 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)
@@ -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)

0 comments on commit e00adba

Please sign in to comment.
You can’t perform that action at this time.