add feedforward for global linear attention as well

lucidrains · Jun 5, 2021 · dacdef6 · dacdef6
1 parent 295f642
commit dacdef6
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/egnn_pytorch/egnn_pytorch.py b/egnn_pytorch/egnn_pytorch.py
@@ -146,13 +146,25 @@ def __init__(
         self.attn1 = Attention(dim, heads, dim_head)
         self.attn2 = Attention(dim, heads, dim_head)
 
+        self.ff = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim)
+        )
+
     def forward(self, x, queries, mask = None):
         res_x, res_queries = x, queries
         x, queries = self.norm_seq(x), self.norm_queries(queries)
 
         induced = self.attn1(queries, x, mask = mask)
         out     = self.attn2(x, induced)
-        return out + res_x, induced + res_queries
+
+        x =  out + res_x
+        queries = induced + res_queries
+
+        x = self.ff(x) + x
+        return x, queries
 
 # classes
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'egnn-pytorch',
   packages = find_packages(),
-  version = '0.2.2',
+  version = '0.2.3',
   license='MIT',
   description = 'E(n)-Equivariant Graph Neural Network - Pytorch',
   author = 'Phil Wang, Eric Alcaide',