attn shape flatten

pyg-team · Jan 3, 2020 · 3284981 · 3284981
1 parent c84cee3
commit 3284981
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/triangles_sag_pool.py b/examples/triangles_sag_pool.py
@@ -13,7 +13,7 @@
 
 class HandleNodeAttention(object):
     def __call__(self, data):
-        data.attn = torch.softmax(data.x, dim=0)
+        data.attn = torch.softmax(data.x, dim=0).flatten()
         data.x = None
         return data