From f86076a31400155bb312cbb193d391045db13c75 Mon Sep 17 00:00:00 2001
From: neulus <softmax@kookmin.ac.kr>
Date: Mon, 29 Sep 2025 16:18:28 +0900
Subject: [PATCH] better arti

---
 main.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index 7b048d0..62ba1cf 100644
--- a/main.py
+++ b/main.py
@@ -1,3 +1,4 @@
+import os
 from typing import cast
 
 import torch
@@ -41,6 +42,11 @@ train_dataset.set_format(type="torch", columns=["x0", "x1"])
 
 wandb.init(project="cloud-removal-kmu")
 
+if not (wandb.run and wandb.run.name):
+    raise Exception("nope")
+
+os.makedirs(f"artifact/{wandb.run.name}", exist_ok=True)
+
 batch_size = 16
 for epoch in range(100):
     lossbin = {i: 0 for i in range(10)}
@@ -77,6 +83,6 @@ for epoch in range(100):
                 "model_state_dict": model.state_dict(),
                 "optimizer_state_dict": optimizer.state_dict(),
             },
-            f"checkpoint_epoch_{epoch + 1}.pt",
+            f"artifact/{wandb.run.name}/checkpoint_epoch_{epoch + 1}.pt",
         )
 wandb.finish()