From e7d36146545f0b664832f3671445aa1afeda9f2c Mon Sep 17 00:00:00 2001
From: patrickgadd <patrickgadd@gmail.com>
Date: Tue, 6 Sep 2022 16:38:17 +0200
Subject: [PATCH] Moves LayerNorm to output of the Encoder's sub-layers

---
 allrank/models/transformer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/allrank/models/transformer.py b/allrank/models/transformer.py
index c09aaa8..dd1b964 100644
--- a/allrank/models/transformer.py
+++ b/allrank/models/transformer.py
@@ -102,8 +102,8 @@ def forward(self, x, sublayer):
         :param sublayer: layer through which to pass the input prior to applying the sum
         :return: output of shape [batch_size, slate_length, output_dim]
         """
-        return x + self.dropout(
-            sublayer(self.norm(x)))
+        return self.norm(x + self.dropout(
+            sublayer(x)))
 
 
 class EncoderLayer(nn.Module):