Update configuration_minGRULM.py

configuration_minGRULM.py

@@ -26,3 +26,22 @@ class MinGRULMConfig(PretrainedConfig):
         self.pad_vocab_size_multiple = pad_vocab_size_multiple
         self.initializer_range = initializer_range
         super().__init__(**kwargs)
+
+    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
+        state_dict['config'] = self.config.__dict__
+    
+        return state_dict
+