give locality to feedforwards even in cross attention block

alphafold2_pytorch/alphafold2.py

@@ -851,7 +851,7 @@ class Alphafold2(nn.Module):
 
             layers.append(nn.ModuleList([
                 intercept_fn(context = False, attn = prenorm_cross(cross_attn_fn())),
-                prenorm(FeedForward(dim = dim, dropout = ff_dropout)),
+                prenorm(LocalFeedForward(dim = dim, hidden_dim = dim * 4, dropout = ff_dropout)),
                 intercept_fn(context = True, attn = prenorm_cross(cross_attn_fn())),
                 prenorm(FeedForward(dim = dim, dropout = ff_dropout)),
             ]))

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'alphafold2-pytorch',
   packages = find_packages(),
-  version = '0.0.89',
+  version = '0.0.90',
   license='MIT',
   description = 'AlphaFold2 - Pytorch',
   author = 'Phil Wang, Eric Alcaide',