Update modeling_minGRULM.py

modeling_minGRULM.py

@@ -61,13 +61,19 @@ class MinGRULMForCausalLM(MinGRULMPreTrainedModel):
 
         # Language modeling head
         self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
-        
-        # Copy weights instead of sharing them
-        with torch.no_grad():
-            self.lm_head.weight.data.copy_(self.model.min_gru_model.token_emb.weight.data)
-        
+
         self.post_init()
 
+    def post_init(self):
+        super().post_init()
+
+        # Ensure tied weights
+        self.tie_weights()
+
+    def tie_weights(self):
+        # Tie lm_head weights to the embedding layer weights
+        self.lm_head.weight = self.model.min_gru_model.token_emb.weight
+
     def get_input_embeddings(self):
         return self.model.min_gru_model.token_emb