Update modeling_super_linear.py

modeling_super_linear.py

@@ -561,7 +561,6 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
 
     def forward(self,
                 inputs_embeds: torch.Tensor = None,   
-                prediction_len: int = None,
                 attention_mask: Optional[torch.Tensor] = None,
                 past_key_values: Optional[Tuple] = None,
                 use_cache: bool = True,
@@ -577,11 +576,11 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
         return CausalLMOutputWithCrossAttentions(loss=None,logits=preds,past_key_values=None,hidden_states=None,attentions=None,)
 
 
-    def prepare_inputs_for_generation(self, inputs_embeds, past_key_values=None, prediction_len=None, **kwargs):
+    def prepare_inputs_for_generation(self, inputs_embeds, past_key_values=None, **kwargs):
         if past_key_values is not None:
             # only feed the last new step
             inputs_embeds = inputs_embeds[:, -1:, :]
-        return {"inputs_embeds": inputs_embeds, "past_key_values": past_key_values, "prediction_len": prediction_len}
+        return {"inputs_embeds": inputs_embeds, "past_key_values": past_key_values}
 
     def _reorder_cache(self, past, beam_idx, **kwargs):
         return past  # backbone keeps no KV cache