fix more mistakes

Author: lucidrains

nGPT_pytorch/nGPT.py

@@ -71,7 +71,7 @@ def __init__(
         self.to_v = NormLinear(dim, dim_inner)
 
         self.rotary_emb = RotaryEmbedding(dim_head)
-        self.qk_scale = nn.Parameter(torch.ones(dim_head) * (dim_head ** -0.25))
+        self.qk_scale = nn.Parameter(torch.ones(dim_head) * (dim_head ** 0.25))
 
         self.norm_qk = norm_qk
         self.split_heads = Rearrange('b n (h d) -> b h n d', h = heads)
@@ -207,10 +207,10 @@ def forward(
 
         for (attn, ff), (attn_alpha, ff_alpha) in zip(self.layers, self.residual_lerp_scales):
 
-            attn_out = attn(tokens)
+            attn_out = l2norm(attn(tokens))
             tokens = l2norm(tokens.lerp(attn_out, attn_alpha))
 
-            ff_out = ff(tokens)
+            ff_out = l2norm(ff(tokens))
             tokens = l2norm(tokens.lerp(ff_out, ff_alpha))
 
         logits = self.to_logits(tokens)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nGPT-pytorch"
-version = "0.0.5"
+version = "0.0.6"
 description = "nGPT"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

train.py

@@ -89,7 +89,7 @@ def base_decoding(
 model = nGPT(
     num_tokens = 256,
     dim = 512,
-    depth = 6
+    depth = 8
 ).to(device)
 
 # prepare enwik8 data