Update modeling_super_linear.py

modeling_super_linear.py

@@ -597,11 +597,14 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
         if len(x.shape) == 2:
             x = x.unsqueeze(1)  # [B, 1, L]
         
-        scale_factor = 512/x.shape[-1]
+        scale_factor = int(np.ceil(512/x.shape[-1]))
         upsample = interpolate(x, scale_factor=scale_factor, mode='linear')  # [B, C, new_L]
         
         # Take last 500 timesteps and rearrange to [B, L, C]
-        upsample = upsample.permute(0, 2, 1)[:, -target_len:, :]
+        if target_len < 512:
+            upsample = upsample.permute(0, 2, 1)
+        else:
+            upsample = upsample.permute(0, 2, 1)[:, -target_len:, :]
         
         # If input was 2D, remove the channel dimension
         if x.shape[1] == 1:
@@ -636,10 +639,12 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
             x_enc = self.upsample_interpolate(x_enc)
             pass
 
-
+        scale_factor = int(np.ceil(512/x.shape[-1]))
         # backbone returns (B, pred_len, C)
+        self.backbone.inf_pred_len = 96*scale_factor
         
         preds = self.backbone(x_enc)
+        x_enc = self.upsample_interpolate(x_enc,x_enc.shape[1)
         #preds = self.fourier_downsample_dim1(preds,96)
         #preds = self.revin_layer(preds, 'denorm')