properly mask for structure module

lucidrains · Aug 23, 2021 · c85682e · c85682e
1 parent dda22b5
commit c85682e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/alphafold2_pytorch/alphafold2.py b/alphafold2_pytorch/alphafold2.py
@@ -337,7 +337,7 @@ def forward(self, x, mask = None):
         if exists(mask):
             # masked mean, if there are padding in the rows of the MSA
             mask = rearrange(mask, 'b m i -> b m i () ()') * rearrange(mask, 'b m j -> b m () j ()')
-            outer = outer.masked_fill(mask, 0.)
+            outer = outer.masked_fill(~mask, 0.)
             outer = outer.mean(dim = 1) / (mask.sum(dim = 1) + self.eps)
         else:
             outer = outer.mean(dim = 1)
@@ -860,6 +860,7 @@ def forward(
 
                 single_repr = self.ipa_block(
                     single_repr,
+                    mask = mask,
                     pairwise_repr = pairwise_repr,
                     rotations = rotations,
                     translations = translations

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'alphafold2-pytorch',
   packages = find_packages(),
-  version = '0.4.26',
+  version = '0.4.29',
   license='MIT',
   description = 'AlphaFold2 - Pytorch',
   author = 'Phil Wang, Eric Alcaide',