minor fix

lean-dojo · Apr 12, 2024 · 4d6b89b · 4d6b89b
1 parent 6aaf3ac
commit 4d6b89b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/generator/model.py b/generator/model.py
@@ -2,6 +2,7 @@
 
 import os
 import torch
+import shutil
 import openai
 import pickle
 from lean_dojo import Pos
@@ -280,7 +281,7 @@ def on_validation_epoch_end(self) -> None:
         logger.info(f"Pass@1: {acc}")
 
         if os.path.exists(ckpt_path):
-            os.remove(ckpt_path)
+            shutil.rmtree(ckpt_path)
 
     ##############
     # Prediction #