Merge pull request #27 from brennanaba/patch-1

Code breaks if there is no mask defined
lucidrains · Apr 5, 2021 · 6d423c9 · 6d423c9
2 parents 876f995 + 881db08
commit 6d423c9
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/egnn_pytorch/egnn_pytorch.py b/egnn_pytorch/egnn_pytorch.py
@@ -240,8 +240,9 @@ def forward(self, feats, coors, edges = None, mask = None, adj_mat = None):
             coors_out = coors
 
         if exists(self.node_mlp):
-            m_ij_mask = rearrange(mask, '... -> ... ()')
-            m_ij = m_ij.masked_fill(~m_ij_mask, 0.)
+            if exists(mask):
+                m_ij_mask = rearrange(mask, '... -> ... ()')
+                m_ij = m_ij.masked_fill(~m_ij_mask, 0.)
 
             if self.m_pool_method == 'mean':
                 if exists(mask):