Merge pull request #711 from allenai/epwalsh/fix-unbound-qkv

fix unbound qkv
allenai · Aug 22, 2024 · 213a639 · 213a639
2 parents 4575d40 + 38bf7ce
commit 213a639
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)