use stable softmax in attention

lucidrains · Aug 9, 2021 · b88e9e3 · b88e9e3
1 parent e6ed6f4
commit b88e9e3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/alphafold2_pytorch/alphafold2.py b/alphafold2_pytorch/alphafold2.py
@@ -161,6 +161,7 @@ def forward(self, x, mask = None, attn_bias = None, context = None, context_mask
 
         # attention
 
+        dots = dots - dots.max(dim = -1, keepdims = True).values
         attn = dots.softmax(dim = -1)
         attn = self.dropout(attn)
 

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