From 4ed744ef166d38dce4faf6241b76518322d9e11f Mon Sep 17 00:00:00 2001
From: Giulia Baldini <44327645+giuliabaldini@users.noreply.github.com>
Date: Fri, 11 Oct 2024 08:21:27 +0200
Subject: [PATCH] Only remove folder in sentenpiece check if it was created
 (#1121)

---
 unsloth/save.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/unsloth/save.py b/unsloth/save.py
index 3ba62b4688..3e82a5bcdd 100644
--- a/unsloth/save.py
+++ b/unsloth/save.py
@@ -101,14 +101,17 @@ def check_if_sentencepiece_model(model, temporary_location = "_unsloth_sentencep
     temp_tokenizer = model._saved_temp_tokenizer
     sentencepiece_model = False
     file_location = os.path.join(temporary_location, temp_tokenizer.name_or_path)
+    created_folder = False
     if not os.path.exists(file_location):
+        created_folder = True
         os.makedirs(file_location)
     pass
     temp_tokenizer.save_pretrained(file_location)
     if os.path.isfile(f"{file_location}/tokenizer.model"):
         sentencepiece_model = True
     pass
-    shutil.rmtree(file_location, ignore_errors = True)
+    if created_folder:
+        shutil.rmtree(file_location, ignore_errors = True)
     return sentencepiece_model
 pass