Update modeling_minGRULM.py

modeling_minGRULM.py

@@ -147,14 +147,13 @@ class MinGRULMForCausalLM(PreTrainedModel):
             # Save the model's state_dict (model weights)
             #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()
+
+            for name, param in self.model.min_gru_model.named_parameters():
+                state_dict[f"model.{name}"] = param
         
-            # Add config as a dictionary (not state_dict, since it is not available)
+            for name, param in self.lm_head.named_parameters():
+                state_dict[f"lm_head.{name}"] = param
+
             state_dict['config'] = self.config.__dict__
             torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin"))