modeling update

modeling_super_linear.py

@@ -5,6 +5,8 @@ import torch, torch.nn as nn, torch.nn.functional as F
 from transformers                          import (PreTrainedModel,GenerationMixin,AutoConfig,AutoModelForCausalLM,)
 from transformers.modeling_outputs         import CausalLMOutputWithCrossAttentions
 from .configuration_super_linear           import SuperLinearConfig
+from .configuration_super_linear_fs        import SuperLinearConfigFS
+from typing import Tuple, Union
 
 
 import math
@@ -544,9 +546,9 @@ class superLinear(nn.Module):
     
 "-------------------------------------------------------------------------------------------------------------------"
 class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
-    config_class = SuperLinearConfig
+    config_class: Tuple[type] = (SuperLinearConfig, SuperLinearConfigFS)
 
-    def __init__(self, config: SuperLinearConfig):
+    def __init__(self, config: Union[SuperLinearConfig, SuperLinearConfigFS]):
         super().__init__(config)