From df0edacc60b502d25ff185d4ff82fe0e436dca6b Mon Sep 17 00:00:00 2001
From: mrq <barry.quiggles@protonmail.com>
Date: Fri, 10 Mar 2023 14:04:07 +0000
Subject: [PATCH] fix the cleanup actually only doing 2 despite requesting more
 than 2, surprised no one has pointed it out

---
 src/utils.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/utils.py b/src/utils.py
index 31bf866..747c3ef 100755
--- a/src/utils.py
+++ b/src/utils.py
@@ -707,8 +707,9 @@ class TrainingState():
 			
 		models = sorted([ int(d[:-8]) for d in os.listdir(f'{self.dataset_dir}/models/') if d[-8:] == "_gpt.pth" ])
 		states = sorted([ int(d[:-6]) for d in os.listdir(f'{self.dataset_dir}/training_state/') if d[-6:] == ".state" ])
-		remove_models = models[:-2]
-		remove_states = states[:-2]
+
+		remove_models = models[:-keep]
+		remove_states = states[:-keep]
 
 		for d in remove_models:
 			path = f'{self.dataset_dir}/models/{d}_gpt.pth'