From 361559a8c262b8732a76468c8925dc71d587f1a6 Mon Sep 17 00:00:00 2001
From: Timur Ganiev <ganiev.tmr@gmail.com>
Date: Sun, 26 Sep 2021 10:24:33 +0300
Subject: [PATCH] Fix layer norm usage (#5)

---
 src/perceiver/attention.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/perceiver/attention.py b/src/perceiver/attention.py
index cea8df3..a9fcdd2 100644
--- a/src/perceiver/attention.py
+++ b/src/perceiver/attention.py
@@ -128,6 +128,7 @@ def __init__(
         """
         super().__init__()
         self.layer_norm = nn.LayerNorm(hidden_dim)
+        self.qkv_layer_norm = nn.LayerNorm(hidden_dim)
         self.attention = MultiHeadAttention(
             kv_dim=hidden_dim,
             q_dim=hidden_dim,
@@ -157,7 +158,7 @@ def forward(
         )
         attention = self.dropout(attention)
         x = x + attention
-        x = x + self.mlp(x)
+        x = x + self.mlp(self.qkv_layer_norm(x))
         return x
 
 
@@ -195,6 +196,7 @@ def __init__(
         self.use_query_residual = use_query_residual
         self.kv_layer_norm = nn.LayerNorm(kv_dim)
         self.q_layer_norm = nn.LayerNorm(q_dim)
+        self.qkv_layer_norm = nn.LayerNorm(q_dim)
         self.attention = MultiHeadAttention(
             kv_dim=kv_dim,
             q_dim=q_dim,
@@ -219,8 +221,8 @@ def forward(
                 in [0, 1]. Defaults to None.
         """
         attention = self.attention(
-            inputs_kv=inputs_kv,
-            inputs_q=inputs_q,
+            inputs_kv=self.kv_layer_norm(inputs_kv),
+            inputs_q=self.q_layer_norm(inputs_q),
             attention_mask=attention_mask
         )
         attention = self.dropout(attention)
@@ -228,5 +230,5 @@ def forward(
             x = inputs_q + attention
         else:
             x = attention
-        x = x + self.mlp(x)
+        x = x + self.mlp(self.qkv_layer_norm(x))
         return x