Update modeling_minGRULM.py

modeling_minGRULM.py

@@ -130,18 +130,22 @@ class MinGRULMForCausalLM(PreTrainedModel):
     def save_pretrained(self, save_directory, safe_serialization: Optional[bool] = True):
         """
         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.
         """
+        # Create the save directory if it doesn't exist
+        os.makedirs(save_directory, exist_ok=True)
+    
+        # Save the model's state_dict (model weights)
+        torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))
+    
+        # Save the configuration
+        self.config.save_pretrained(save_directory)
+    
+        # Optionally print messages based on the safe_serialization flag
         if safe_serialization:
-            print("true")
-            # Save the model's state_dict (model weights)
-            torch.save(self.state_dict(), f"{save_directory}/pytorch_model.bin")
-            # Save the configuration to a file
-            self.config.save_pretrained(save_directory)
+            print("Model and configuration have been saved safely.")
         else:
-            print("false")
-            # If not safe_serialization, you can use the default save mechanism from the base class
-            super().save_pretrained(save_directory)
+            print("Model and configuration have been saved (unsafe serialization).")