From 4d6b89b3efd8b55c280fb41b47843fddc586f764 Mon Sep 17 00:00:00 2001
From: Kaiyu Yang <kaiyuy@caltech.edu>
Date: Fri, 12 Apr 2024 04:21:23 +0000
Subject: [PATCH] minor fix

---
 generator/model.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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 #