add extra feedforward to perceiver io

lucidrains · Aug 29, 2021 · 858aab3 · 858aab3
1 parent 3b70ebe
commit 858aab3
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/perceiver_pytorch/perceiver_io.py b/perceiver_pytorch/perceiver_io.py
@@ -149,6 +149,8 @@ def __init__(
             ]))
 
         self.decoder_cross_attn = PreNorm(queries_dim, Attention(queries_dim, latent_dim, heads = cross_heads, dim_head = cross_dim_head), context_dim = latent_dim)
+        self.decoder_ff = PreNorm(queries_dim, FeedForward(queries_dim))
+
         self.to_logits = nn.Linear(queries_dim, logits_dim) if exists(logits_dim) else nn.Identity()
 
     def forward(
@@ -177,6 +179,7 @@ def forward(
         # cross attend from decoder queries to latents
 
         latents = self.decoder_cross_attn(queries, context = x)
+        latents = self.decoder_ff(latents)
 
         # final linear out
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'perceiver-pytorch',
   packages = find_packages(),
-  version = '0.5.1',
+  version = '0.6.0',
   license='MIT',
   description = 'Perceiver - Pytorch',
   author = 'Phil Wang',