clean up

pyg-team · Jul 6, 2019 · 3935ccf · 3935ccf
1 parent 29e36da
commit 3935ccf
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 28 deletions.
diff --git a/torch_geometric/nn/pool/topk_pool.py b/torch_geometric/nn/pool/topk_pool.py
@@ -9,12 +9,11 @@
 
 def topk(x, ratio, batch, min_score=None, tol=1e-7):
     if min_score is not None:
-        # make sure that we do not drop all nodes in a graph
+        # Make sure that we do not drop all nodes in a graph.
         scores_max = scatter_max(x, batch)[0][batch] - tol
-        scores_min = torch.min(scores_max, x.new_full((1, ), min_score))
+        scores_min = scores_max.clamp(max=min_score)
 
         perm = torch.nonzero(x > scores_min).view(-1)
-
     else:
         num_nodes = scatter_add(batch.new_ones(x.size(0)), batch, dim=0)
         batch_size, max_num_nodes = num_nodes.size(0), num_nodes.max().item()

diff --git a/torch_geometric/transforms/handle_node_attention.py b/torch_geometric/transforms/handle_node_attention.py