Merge branch 'main' into version-dolma-flan-change

allenai · Aug 26, 2024 · 9c82609 · 9c82609
2 parents bd319be + 213a639
commit 9c82609
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/olmo/model.py b/olmo/model.py
@@ -736,11 +736,13 @@ def forward(
         # apply norm before
         if not self.config.norm_after:
             if self._activation_checkpoint_fn is not None:
-                qkv = self._activation_checkpoint_fn(self.attn_norm, x)
+                h = self._activation_checkpoint_fn(self.attn_norm, x)
             else:
-                qkv = self.attn_norm(x)
+                h = self.attn_norm(x)
+        else:
+            h = x
 
-        qkv = self.att_proj(qkv)
+        qkv = self.att_proj(h)
 
         if self.config.clip_qkv is not None:
             qkv.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)