fix

stanford-crfm · Feb 16, 2024 · 4136e66 · 4136e66
1 parent 446fb62
commit 4136e66
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/models/gpt2.py b/src/levanter/models/gpt2.py
@@ -281,7 +281,7 @@ def true_fun(_):
                 # not into ff for now
                 # Operations if layer_idx equals 4
                 # sum over sequence length
-                return hax.sin(0.1*hax.sum(prev_x, axis='position')) + attn_output + ff_output
+                return hax.sin(*hax.sum(prev_x, axis='position')) + attn_output + ff_output
             def false_fun(_):
                 # Otherwise return the same tensor
                 # as expected