Update modeling_mos_mamba.py

modeling_mos_mamba.py

@@ -875,8 +875,8 @@ class MoSMambaForCausalLM(MoSMambaPreTrainedModel):
         super().__init__(config)
         self.backbone = MoSMambaModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.num_selectivities = 6
-        self.num_selectivities_per_tok = 2
+        self.num_selectivities = config.num_selectivities
+        self.num_selectivities_per_tok = config.num_selectivities_per_tok
         self.router_aux_loss_coef = 0.02
         # Initialize weights and apply final processing
         self.post_init()