diff --git a/generator/model.py b/generator/model.py
index 65f47cc..528a4e4 100644
--- 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 #