Update cumo/model/language_model/llava_llama.py

cumo/model/language_model/llava_llama.py

@@ -72,22 +72,24 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
 
         if inputs_embeds is None:
             (
-                input_ids,
+                inputs,
                 position_ids,
                 attention_mask,
-                past_key_values,
+                _,
                 inputs_embeds,
-                labels
+                _,
+                *_
             ) = self.prepare_inputs_labels_for_multimodal(
-                input_ids,
+                inputs,
                 position_ids,
                 attention_mask,
-                past_key_values,
-                labels,
+                None,
+                None,
                 images,
-                image_sizes
+                image_sizes=image_sizes
             )
 
+
         return super().forward(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -121,7 +123,8 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
                 attention_mask,
                 _,
                 inputs_embeds,
-                _
+                _,
+                *_
             ) = self.prepare_inputs_labels_for_multimodal(
                 inputs,
                 position_ids,
@@ -131,6 +134,7 @@ class LlavaLlamaForCausalLM(LlamaForCausalLM, LlavaMetaForCausalLM):
                 images,
                 image_sizes=image_sizes
             )
+
         else:
             inputs_embeds = self.get_model().embed_tokens(inputs)