Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -1205,6 +1205,14 @@ class GemmoeForCausalLM(GemmoePreTrainedModel):
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+        attention_mask = attention_mask.to(device) if attention_mask is not None else None
+        position_ids = position_ids.to(device) if position_ids is not None else None
+        past_key_values = [t.to(device) for t in past_key_values] if past_key_values is not None else None
+        inputs_embeds = inputs_embeds.to(device) if inputs_embeds is not None else None
+        labels = labels.to(device) if labels is not None else None
+        cache_position = cache_position.to(device) if cache_position is not None else None
+
         outputs = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,