Update modeling_mos_mamba.py

modeling_mos_mamba.py

@@ -691,6 +691,8 @@ class MoSMambaPreTrainedModel(PreTrainedModel):
             elif self.config.time_step_init_scheme == "random":
                 nn.init.uniform_(module.dt_proj.weight, -dt_init_std, dt_init_std)
 
+            nn.init.xavier_uniform_(module.gate.weight, gain=0.1)
+
             dt = torch.exp(
                 torch.rand(self.config.intermediate_size)
                 * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min))