From 03bd6ad9d8899cf891850ba7d1dd3e63107f1a34 Mon Sep 17 00:00:00 2001
From: Filip Stefaniuk <filip@desktop-ctgf7km.home>
Date: Sat, 7 Sep 2024 08:09:33 -0400
Subject: [PATCH] Add models cleanup to training script

---
 scripts/train.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/scripts/train.py b/scripts/train.py
index 020bfcd..8ee088d 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -251,7 +251,11 @@ def main():
         batch_size=batch_size, train=False, num_workers=3),
         ckpt_path=ckpt_path)
 
-    # TODO: Clean up non-best models to save space on wandb
+    # Clean up models that do not have best/latest tags, to save space on wandb
+    for artifact in wandb.Api().run(run.path).logged_artifacts():
+        if artifact.type == "model" and not artifact.aliases:
+            logging.info(f"Deleting artifact {artifact.name}")
+            artifact.delete()
 
 
 if __name__ == '__main__':