Update modeling_super_linear.py

modeling_super_linear.py

@@ -213,26 +213,23 @@ class RLinear(nn.Module):
     def forward(self, x):
         # x: [Batch, Input length,Channel]
         x_shape = x.shape
-        print(x.shape)
-        if len(x_shape) == 2:
-            x = x.unsqueeze(-1)
-
-        B,L,V    = x.shape
-        if L <  self.seq_len and self.zero_shot_Linear is None:
-            print(F"New Lookback :{L}")
-            self.transform_model(L)
-
-        if L <  self.seq_len:
+        if x.shape[1] < self.seq_len:
+            if self.zero_shot_Linear is None:
+                self.transform_model(L)
             x = x.clone()
             x = self.revin_layer(x, 'norm')
             x = F.linear(x, self.zero_shot_Linear)
             x = self.revin_layer(x, 'denorm')
+            return x
 
-        else:
-            x = x.clone()
-            x = self.revin_layer(x, 'norm')
-            x = self.Linear(x.permute(0,2,1)).permute(0,2,1).clone()
-            x = self.revin_layer(x, 'denorm')
+                
+        if len(x_shape) == 2:
+            x = x.unsqueeze(-1)
+
+        x = x.clone()
+        x = self.revin_layer(x, 'norm')
+        x = self.Linear(x.permute(0,2,1)).permute(0,2,1).clone()
+        x = self.revin_layer(x, 'denorm')
 
         if len(x_shape) == 2:
             x = x.squeeze(-1)