Update modeling_super_linear.py

modeling_super_linear.py

@@ -525,9 +525,10 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
   
 
         # the backbone keeps its own Config dataclass, so build one on‑the‑fly:
-        backbone_cfg   = type("Cfg", (), config.to_dict())()
-        self.args      = backbone_cfg
-        self.backbone  = superLinear(backbone_cfg)
+        backbone_cfg     = type("Cfg", (), config.to_dict())()
+        self.args        = backbone_cfg
+        self.backbone    = superLinear(backbone_cfg)
+        self.revin_layer = RevIN(num_features = None, affine=False, norm_type = None, subtract_last = False)
         self.post_init()                             
 
     # ------------------------------------------------------------------
@@ -589,15 +590,17 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
         
         # backbone expects (B, C, L)
         x_enc = inputs_embeds
+                
+
         
         if x_enc.shape[1] < 512:
             x_enc = self.fourier_interp_dim1(x_enc)
-            mean = x_enc.mean()
-            std  = x_enc.std().clamp_min(1e-6)
-            x_enc = (x_enc - mean) / std
+            x_enc = self.revin_layer(x_enc, 'norm')
+
    
         # backbone returns (B, pred_len, C)
         preds = self.backbone(x_enc)
+        preds = self.revin_layer(preds, 'denorm')
         return CausalLMOutputWithCrossAttentions(loss=None,logits=preds,past_key_values=None,hidden_states=None,attentions=None,)