Update modeling_super_linear.py

modeling_super_linear.py

@@ -593,22 +593,21 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
         return y
     
     def upsample_interpolate(self, x,scale_factor, target_len: int = 512):
-        # Add channel dimension if input is 2D
-        size = len(x.shape)
-        if size == 2:
-            x = x.unsqueeze(1)  
-        
-        print(target_len)
-        upsample =  F.interpolate(x, size=target_len, mode='linear', align_corners=False) 
+        was_2d = x.dim() == 2                       
+    
+        if was_2d:                                  # [B, L] -> [B, 1, L]
+            x = x.unsqueeze(1)
+        else:                                       # [B, L, C] -> [B, C, L]
+            x = x.permute(0, 2, 1)
+    
         
-        # If input was 2D, remove the channel dimension
-        if size == 2:
-            upsample = upsample.squeeze(1)
-
-        upsample = upsample.float()
-            
-        #print(f"Upsampled shape: {upsample.shape}")
-        return upsample
+        x_up = F.interpolate(x, size=target_len, mode='linear', align_corners=False)                                            
+    
+        # ── restore original layout ─────────────────────────────────────────────────
+        if was_2d:                                  # back to [B, target_len]
+            return x_up.squeeze(1).float()
+        else:                                       # back to [B, target_len, C]
+            return x_up.permute(0, 2, 1).float()