Update modeling_super_linear.py

modeling_super_linear.py

@@ -199,14 +199,15 @@ 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):
-        W              = self.Linear.weight.detach()
-        new_W          = W[:, -new_lookback:] 
-        original_norm  = torch.norm(W, p=2)
-        new_norm       = torch.norm(new_W, p=2)
-        final_scaling  = original_norm / new_norm if new_norm.item() != 0 else 1.0
-        #final_scaling  = 1
-        new_W          = new_W * final_scaling
+    def transform_model(self,new_lookback,mode):
+        if mode == 2: 
+            W              = self.Linear.weight.detach()
+            new_W          = W[:, -new_lookback:] 
+            original_norm  = torch.norm(W, p=2)
+            new_norm       = torch.norm(new_W, p=2)
+            final_scaling  = original_norm / new_norm if new_norm.item() != 0 else 1.0
+            #final_scaling  = 1
+            new_W          = new_W * final_scaling
 
         self.zero_shot_Linear        = new_W
 
@@ -219,13 +220,12 @@ class RLinear(nn.Module):
                 #print(F"new Lookkback : {x.shape[1]}")
                 self.transform_model(x.shape[1])
 
-
             x = x.clone()
-            x = x * (x.shape[1]/512)
+            #x = x * (x.shape[1]/512)
             x = self.revin_layer(x, 'norm')
             x = F.linear(x, self.zero_shot_Linear)
             x = self.revin_layer(x, 'denorm')
-            x = x * (512/x.shape[1])
+            #x = x * (512/x.shape[1])
             return x