Update modeling_minGRULM.py

modeling_minGRULM.py

@@ -129,13 +129,14 @@ class MinGRULMForCausalLM(PreTrainedModel):
         model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
         return model
 
-    def save_pretrained(self, save_directory, safe_serialization: Optional[bool] = True):
+    def save_pretrained(self, save_directory, safe_serialization: Optional[bool] = True, **kwargs):
         """
         Save the model and configuration to a directory.
     
         Args:
             save_directory (str): Directory to save the model.
             safe_serialization (bool, optional): Whether to use safe serialization. Defaults to True.
+            kwargs: Additional arguments like max_shard_size (ignored in this implementation).
         """
         import os
         os.makedirs(save_directory, exist_ok=True)
@@ -144,13 +145,13 @@ class MinGRULMForCausalLM(PreTrainedModel):
             print("Saving with safe serialization.")
             
             state_dict = {}
-
+    
             for name, param in self.model.min_gru_model.named_parameters():
                 state_dict[f"model.{name}"] = param
-        
-            for name, param in self.lm_head.named_parameters():
-                state_dict[f"lm_head.{name}"] = param
-
+    
+            for name, param in self.classifier.named_parameters():
+                state_dict[f"classifier.{name}"] = param
+    
             state_dict['config'] = self.config.__dict__
             torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin"))