Update modeling_super_linear.py

modeling_super_linear.py

@@ -592,10 +592,23 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
     
         return y
     
-    def upsample_interpolate(self,x, target_len: int = 512):
-        scale_factor = 512/x.shape[1]
-        upsample     = interpolate(x, scale_factor=scale_factor, mode='linear').permute(0,2,1)[:, -500:, :]
-        print(upsample.shape)
+    def upsample_interpolate(self, x, 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 = 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)[:, -500:, :]
+        
+        # If input was 2D, remove the channel dimension
+        if x.shape[1] == 1:
+            upsample = upsample.squeeze(-1)
+            
+        print(f"Upsampled shape: {upsample.shape}")
+        return upsample