Update modeling_super_linear.py

modeling_super_linear.py

@@ -594,20 +594,15 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
     
     def upsample_interpolate(self, x,scale_factor, target_len: int = 512):
         # Add channel dimension if input is 2D
-        if len(x.shape) == 2:
+        size = len(x.shape)
+        if size == 2:
             x = x.unsqueeze(1)  # [B, 1, L]
         
         
-        upsample = interpolate(x,size =target_len,mode='linear')  # [B, C, new_L]
-        
-        # Take last 500 timesteps and rearrange to [B, L, C]
-        if target_len < 512:
-            upsample = upsample.permute(0, 2, 1)
-        else:
-            upsample = upsample.permute(0, 2, 1)
+        upsample = = F.interpolate(x, size=target_len, mode='linear', align_corners=False) 
         
         # If input was 2D, remove the channel dimension
-        if x.shape[1] == 1:
+        if size == 2:
             upsample = upsample.squeeze(-1)
 
         upsample = upsample.float()