Update starvector_arch.py

starvector_arch.py

@@ -159,23 +159,22 @@ class StarVectorForCausalLM(PreTrainedModel):
         if hasattr(self.model, 'svg_transformer') and hasattr(self.model.svg_transformer, 'gradient_checkpointing_enable'):
             self.model.svg_transformer.gradient_checkpointing_enable()
 
-    def forward(self,  inputs_embeds, input_ids, num_generations, num_logits_to_keep) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
+    def forward(self,  vision_embeds, input_ids, num_generations, num_logits_to_keep) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
         r"""
         Wrapper for the forward pass of the model.
         """
-        device = inputs_embeds.device
+        device = vision_embeds.device
 
         completion_embeds = self.model._get_embeddings(input_ids)
-        inputs_embeds = torch.cat([inputs_embeds.repeat(num_generations, 1, 1), completion_embeds], dim=1)
-        attention_mask = torch.ones_like(inputs_embeds[:, :, 0]).to(device)
+        vision_embeds = torch.cat([vision_embeds.repeat(num_generations, 1, 1), completion_embeds], dim=1)
+        attention_mask = torch.ones_like(vision_embeds[:, :, 0]).to(device)
 
         transformer_outputs = self.model.svg_transformer.transformer.transformer(
-            inputs_embeds=inputs_embeds, 
+            inputs_embeds=vision_embeds, 
             attention_mask=attention_mask,
         )
         hidden_states = transformer_outputs[0]
 
-        # If GRPO requested only the last tokens, slice accordingly.
         if num_logits_to_keep > 0:
             lm_logits = self.model.svg_transformer.transformer.lm_head(hidden_states[:, -num_logits_to_keep:, :])
         else: