fix neighbor calculation bug in the presence of masking

lucidrains · May 20, 2021 · 24872cd · 24872cd
1 parent 8265086
commit 24872cd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/egnn_pytorch/egnn_pytorch.py b/egnn_pytorch/egnn_pytorch.py
@@ -190,7 +190,7 @@ def forward(self, feats, coors, edges = None, mask = None, adj_mat = None):
             ranking = rel_dist[..., 0].clone()
 
             if exists(mask):
-                rank_mask = mask[:, None, :] * mask[:, None, :]
+                rank_mask = mask[:, :, None] * mask[:, None, :]
                 ranking.masked_fill_(~rank_mask, 1e5)
 
             if exists(adj_mat):

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'egnn-pytorch',
   packages = find_packages(),
-  version = '0.1.11',
+  version = '0.1.12',
   license='MIT',
   description = 'E(n)-Equivariant Graph Neural Network - Pytorch',
   author = 'Phil Wang, Eric Alcaide',