Update cumo/model/llava_arch.py

cumo/model/llava_arch.py

@@ -160,7 +160,7 @@ class LlavaMetaForCausalLM(ABC):
         mlp_router_z_loss = None
         vision_tower = self.get_vision_tower()
         if vision_tower is None or images is None or input_ids.shape[1] == 1:
-            return input_ids, position_ids, attention_mask, past_key_values, None, labels, clip_balanced_loss, clip_router_z_loss, mlp_balanced_loss, mlp_router_z_loss 
+            return input_ids, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels 
         
         if type(images) is list or images.ndim == 5:
             if type(images) is list: