Update modeling_quiet.py

modeling_quiet.py

@@ -1283,6 +1283,11 @@ class QuietForCausalLM(QuietPreTrainedModel):
                 self.talk_head = nn.ModuleList([nn.Sequential(
                     nn.Linear(talk_input_dim, talk_output_dim, bias=False)
                 )])
+                
+        self.apply(self._init_weights)
+
+        # Add dropout regularization
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -1304,6 +1309,14 @@ class QuietForCausalLM(QuietPreTrainedModel):
 
     def get_decoder(self):
         return self.model
+    
+    def _init_weights(self, module):
+    if isinstance(module, nn.Linear):
+        nn.init.xavier_uniform_(module.weight)
+        if module.bias is not None:
+            nn.init.constant_(module.bias, 0)
+    elif isinstance(module, nn.Embedding):
+        nn.init.xavier_uniform_(module.weight)
 
     @torch.no_grad()
     def infer(