Update modeling_super_linear.py

modeling_super_linear.py

@@ -561,6 +561,7 @@ 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,
@@ -572,18 +573,17 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
             raise ValueError("Pass the time‑series as `inputs_embeds`")
         
         # backbone expects (B, C, L)
-        preds = self.backbone(inputs_embeds)
+        preds = self.backbone(inputs_embeds, inf_pred_len=prediction_len)
         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, **kwargs):
+    def prepare_inputs_for_generation(self, inputs_embeds, past_key_values=None, prediction_len=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}
+        return {"inputs_embeds": inputs_embeds, "past_key_values": past_key_values, "prediction_len": prediction_len}
 
     def _reorder_cache(self, past, beam_idx, **kwargs):
         return past  # backbone keeps no KV cache
 
 
-