Update modeling_super_linear.py

modeling_super_linear.py

@@ -200,7 +200,7 @@ class RLinear(nn.Module):
         self.revin_layer         = RevIN(num_features = None, affine=False, norm_type = None, subtract_last = False)
         self.zero_shot_Linear    = None
 
-    def transform_model(self,new_lookback,mode):
+    def transform_model(self,x,new_lookback,mode):
         if mode == 1:
             W              = self.Linear.weight.detach()
             new_W          = W[:, -new_lookback:] 
@@ -226,25 +226,19 @@ class RLinear(nn.Module):
         else:
             W = self.Linear.weight.detach()              # (out_features, seq_len)
             
-            # 1️⃣ keep the last `new_lookback` columns
-            W_tail = W[:, -new_lookback:]                # (out_features, new_lookback)
-            
-            # 2️⃣ resize those columns back to `seq_len`
-            #     ─ use 4-D shape (N=1, C=1, H=out, W=new_lookback)
-            W4d = W_tail.unsqueeze(0).unsqueeze(0)        
-            
+            W4d = W.unsqueeze(0).unsqueeze(0)            # (1, 1, out, in)
+
+            # resize   H → self.horizon   and   W → new_lookback
             new_W = F.interpolate(
                         W4d,
-                        size=(new_lookback, self.seq_len),  # keep rows = out_features,
-                                                            #      stretch cols to seq_len
+                        size=(self.horizon, new_lookback),   # (H_out, W_out)
                         mode='bilinear',
                         align_corners=False
-                    )[0, 0]                                # → (out_features, seq_len)
-            
-            # 3️⃣ concatenate on the column axis
-            W_now = torch.cat((W_tail, new_W), dim=1)      # (out_features, new_lookback + seq_len)
-            
-            self.zero_shot_Linear = W_now
+                    )[0, 0]                                 # drop the two singleton dims
+
+            x = F.linear(x, new_W)
+            self.zero_shot_Linear = W
+            return x
 
             
 
@@ -254,7 +248,8 @@ class RLinear(nn.Module):
         if x.shape[1] < self.seq_len:
             #if self.zero_shot_Linear is None:
                 #print(F"new Lookkback : {x.shape[1]}")
-            self.transform_model(x.shape[1],3)
+            
+            x = self.transform_model(x,x.shape[1],3)
 
             x = x.clone()
             #x = x * (x.shape[1]/512)