Update modeling_super_linear.py

modeling_super_linear.py

@@ -216,11 +216,17 @@ class RLinear(nn.Module):
         x_shape = x.shape
         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])
+
+            if len(x_shape) == 2:
+                x = x.unsqueeze(-1)
             x = x.clone()
             x = self.revin_layer(x, 'norm')
-            x = F.linear(x, self.zero_shot_Linear).clone()
+            x = F.linear(x.permute(0,2,1), self.zero_shot_Linear).permute(0,2,1).clone()
             x = self.revin_layer(x, 'denorm')
+            if len(x_shape) == 2:
+                x = x.squeeze(-1)
             return x