Refactor Logits Naming

modeling_moonshot_kimia.py

@@ -902,15 +902,15 @@ class MoonshotKimiaForCausalLM(Qwen2PreTrainedModel):
         else:
             hidden_states, mimo_hidden_states = outputs[0], outputs[1]
 
-        audio_logits = self.lm_head(hidden_states)
-        text_logits = self.mimo_output(mimo_hidden_states)
+        text_logits = self.lm_head(hidden_states)
+        audio_logits = self.mimo_output(mimo_hidden_states)
 
         if not return_dict:
-            output = (text_logits, audio_logits) + outputs[2:]
+            output = (audio_logits, text_logits) + outputs[2:]
             return output
         return CausalLMOutputWithPast(
             loss=None,
-            logits=(text_logits, audio_logits),
+            logits=(audio_logits, text_logits),
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,