pre-layernorm for attention

lucidrains · Apr 4, 2022 · 7b46b01 · 7b46b01
1 parent 42a4c51
commit 7b46b01
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/palm_pytorch/palm_pytorch.py b/palm_pytorch/palm_pytorch.py
@@ -106,6 +106,10 @@ def forward(self, x):
         n, device, h = x.shape[1], x.device, self.heads
         q, k, v = self.to_qkv(x).chunk(3, dim = -1)
 
+        # pre layernorm
+
+        x = self.norm(x)
+
         # split heads
 
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), (q, k, v))

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'PaLM-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.1',
+  version = '0.0.2',
   license='MIT',
   description = 'PaLM: Scaling Language Modeling with Pathways - Pytorch',
   author = 'Phil Wang',