[CLEANUP]

Author: Kye Gomez

bitnet/bit_transformer.py

@@ -75,12 +75,15 @@ def __init__(
                     dropout=0.1,
                 ),
             )
+            
+        # Norm
+        self.norm = nn.LayerNorm(dim)
 
     def forward(self, x: Tensor, *args, **kwargs) -> Tensor:
         skip = x
         for attn, ffn in zip(self.layers, self.ffn_layers):
             x, _ = attn(x, x, x, is_causal=True, *args, **kwargs)
-            x = x + skip
+            x = self.norm(x + skip)
             x = ffn(x) + x
         return x
 
@@ -117,8 +120,8 @@ def __init__(
         dim: int,
         depth: int,
         num_tokens: int,
-        heads=8,
-        ff_mult=4,
+        heads: int = 8,
+        ff_mult: int = 4,
     ):
         super().__init__()
         self.emb = nn.Embedding(num_tokens, dim)
@@ -132,8 +135,14 @@ def __init__(
             dim,
             vocab_size=num_tokens,
         )
+        
+        # Norm
+        self.norm = nn.LayerNorm(dim)
 
     def forward(self, x):
         x = self.emb(x)
+        # Post emb norm
+        x = self.norm(x)
+        
         x = self.transformer(x)
         return self.to_logits(x)