Update src/transformers/models/zamba/modeling_zamba.py

Co-authored-by: Arthur <[email protected]>
Zyphra · Sep 24, 2024 · 66b72c8 · 66b72c8
1 parent 4b0fb52
commit 66b72c8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/models/zamba/modeling_zamba.py b/src/transformers/models/zamba/modeling_zamba.py
@@ -438,7 +438,7 @@ def forward(
         query_states = query_states.transpose(1, 2)
         key_states = key_states.transpose(1, 2)
         value_states = value_states.transpose(1, 2)
-        softmax_scale = 1 / (query_states.shape[-1] / 2) ** 0.5
+        softmax_scale = 1 / math.sqrt(self.head_dim / 2)
 
         attn_output = _flash_attention_forward(
             query_states,