Update modeling_minGRULM.py

modeling_minGRULM.py

@@ -121,24 +121,6 @@ class MinGRULMForCausalLM(PreTrainedModel):
             logits=logits,
         )
 
-    def state_dict(self):
-        """
-        Custom state_dict function to return the model's state dict.
-        This includes the wrapped model and any extra components like the language model head.
-        """
-        state_dict = {}
-    
-        # Add min_gru_model's state_dict
-        state_dict['model'] = self.model.min_gru_model.state_dict()
-        
-        # Add lm_head's state_dict
-        state_dict['lm_head'] = self.lm_head.state_dict()
-    
-        # Add config as a dictionary (not state_dict, since it is not available)
-        state_dict['config'] = self.config.__dict__
-        
-        return state_dict
-
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
         """
@@ -163,7 +145,17 @@ class MinGRULMForCausalLM(PreTrainedModel):
             print("Saving with safe serialization.")
             
             # Save the model's state_dict (model weights)
-            state_dict = self.state_dict()
+            #state_dict = self.state_dict()
+            state_dict = {}
+    
+            # Add min_gru_model's state_dict
+            state_dict['model'] = self.model.min_gru_model.state_dict()
+            
+            # Add lm_head's state_dict
+            state_dict['lm_head'] = self.lm_head.state_dict()
+        
+            # Add config as a dictionary (not state_dict, since it is not available)
+            state_dict['config'] = self.config.__dict__
             torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin"))
             
             # Save the configuration