Update modeling_minGRULM.py

modeling_minGRULM.py

@@ -7,6 +7,7 @@ from typing import Optional
 from .configuration_minGRULM import MinGRULMConfig
 from minGRU_pytorch.minGRULM import minGRULM
 
+
 # Wrapper class for device compatibility
 class MinGRULMWrapped(nn.Module):
     def __init__(self, min_gru_model):
@@ -61,9 +62,6 @@ class MinGRULMForCausalLM(MinGRULMPreTrainedModel):
         # Language modeling head
         self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
 
-        # Initialize weights (if required for missing layers)
-        self.initialize_layers()
-
         self.post_init()
 
     def get_input_embeddings(self):
@@ -105,27 +103,4 @@ class MinGRULMForCausalLM(MinGRULMPreTrainedModel):
         return CausalLMOutputWithPast(
             loss=loss,
             logits=logits,
-        )
-
-    def initialize_layers(self):
-        """
-        Initialize missing layers in the model, such as custom layers or parts of the minGRULM.
-        If layers are already initialized, we can skip them.
-        """
-        # Example: Initialize layers manually if needed
-        for name, module in self.model.min_gru_model.named_children():
-            if isinstance(module, nn.Module):
-                if 'token_emb' in name:
-                    # Token embeddings, if needed, you can initialize with a custom scheme
-                    nn.init.xavier_uniform_(module.weight)
-                elif isinstance(module, nn.Linear):
-                    # Initialize Linear layers if not initialized already
-                    if module.weight is not None:
-                        nn.init.xavier_uniform_(module.weight)
-                    if module.bias is not None:
-                        nn.init.zeros_(module.bias)
-                # Initialize other layers similarly, depending on the type
-                elif isinstance(module, nn.LayerNorm):
-                    # Initialize LayerNorm layers
-                    nn.init.constant_(module.weight, 1.0)
-                    nn.init.constant_(module.bias, 0)
+        )