Update gemma2.py

unslothai · Jul 31, 2024 · 8db7e80 · 8db7e80
1 parent cf1054c
commit 8db7e80
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py
@@ -141,7 +141,7 @@ def Gemma2Attention_fast_forward(
 
         # FA uses 1 / sqrt for softmax_scale!
         if not hasattr(self, "_flash_attention_softmax_scale"):
-            self._flash_attention_softmax_scale = 1.0 / self.config.query_pre_attn_scalar**0.5
+            self._flash_attention_softmax_scale = 1.0 / (self.config.query_pre_attn_scalar**0.5)
         pass
 
         Q = Q.transpose(1, 2)