remove 2 redundant normalizations

allenai · Nov 19, 2024 · 206cfc9 · 206cfc9
1 parent a24e58c
commit 206cfc9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/olmo_core/nn/transformer/block.py b/src/olmo_core/nn/transformer/block.py
@@ -303,7 +303,7 @@ def forward(
     ) -> torch.Tensor:
         h = l2_normalize(
             torch.lerp(
-                l2_normalize(x),
+                x,
                 l2_normalize(self.attention(x, max_doc_len=max_doc_len, cu_doc_lens=cu_doc_lens)),
                 (
                     self.attn_alpha * (self.attn_alpha_init_value / self.attn_alpha_init_scaling)
@@ -313,7 +313,7 @@ def forward(
 
         return l2_normalize(
             torch.lerp(
-                l2_normalize(h),
+                h,
                 l2_normalize(self.feed_forward(h)),
                 (self.mlp_alpha * (self.mlp_alpha_init_value / self.mlp_alpha_init_scaling)).abs(),
             )