Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -629,6 +629,9 @@ class GemmoeSdpaAttention(GemmoeAttention):
             key_states = key_states.contiguous()
             value_states = value_states.contiguous()
 
+            
+        causal_mask = causal_mask.to(query_states.dtype)
+
         attn_output = torch.nn.functional.scaled_dot_product_attention(
             query_states,
             key_states,