add final ckpt

2025-09-29 16:19:24 +09:00
parent f86076a314
commit 02ac62fb1d
1 changed files with 9 additions and 0 deletions
--- a/main.py
+++ b/main.py
@@ -85,4 +85,13 @@ for epoch in range(100):
            },
            f"artifact/{wandb.run.name}/checkpoint_epoch_{epoch + 1}.pt",
        )
+
+torch.save(
+    {
+        "epoch": 100,
+        "model_state_dict": model.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+    },
+    f"artifact/{wandb.run.name}/checkpoint_final.pt",
+)
 wandb.finish()