Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -1215,6 +1215,10 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
         )
 
         hidden_states = outputs[0]
+
+        # Ensure hidden_states and lm_head have compatible dtypes
+        hidden_states = hidden_states.to(dtype=self.lm_head.weight.dtype)
+
         logits = self.lm_head(hidden_states)
         logits = logits.float()
 
@@ -1332,6 +1336,7 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
                 tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
             )
         return reordered_past
+    
 @add_start_docstrings(
     """
     The Gemmoe Model transformer with a sequence classification head on top (linear layer).