Update modeling_minGRU.py

modeling_minGRU.py

@@ -109,13 +109,14 @@ class MinGRUForSequenceClassification(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)
@@ -124,13 +125,13 @@ class MinGRUForSequenceClassification(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.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"))
             
@@ -138,3 +139,4 @@ class MinGRUForSequenceClassification(PreTrainedModel):
         else:
             print("Saving without safe serialization.")
             super().save_pretrained(save_directory)
+