Update modeling_super_linear.py

modeling_super_linear.py

@@ -576,7 +576,7 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
         return y
 
 
-    def fourier_downsample_dim1(self,x, target_len: int):
+    def fourier_downsample_dim1(self,x,target_len: int):
     
         L = x.size(1)
         # 1. Forward real FFT along dim-1
@@ -592,12 +592,12 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
     
         return y
     
-    def upsample_interpolate(self, x, target_len: int = 512):
+    def upsample_interpolate(self, x,scale_factor, target_len: int = 512):
         # Add channel dimension if input is 2D
         if len(x.shape) == 2:
             x = x.unsqueeze(1)  # [B, 1, L]
         
-        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]
@@ -636,15 +636,16 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
         
         if x_enc.shape[1] < 512:
             #x_enc = self.revin_layer(x_enc, 'norm')
-            x_enc = self.upsample_interpolate(x_enc)
+            scale_factor = int(np.ceil(512/x_enc.shape[-1]))
+            x_enc = self.upsample_interpolate(x_enc,scale_factor,512)
             pass
 
-        scale_factor = int(np.ceil(512/x_enc.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])
+        x_enc = self.upsample_interpolate(x_enc,1/scale_factor,x_enc.shape[1])
         #preds = self.fourier_downsample_dim1(preds,96)
         #preds = self.revin_layer(preds, 'denorm')