diff --git a/alphafold2_pytorch/alphafold2.py b/alphafold2_pytorch/alphafold2.py
index 29477d3..25b235c 100644
--- 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
index abe60d2..8f7cf21 100644
--- 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',