Add removal of temporal file

optuna · Oct 25, 2024 · 3fc41e8 · 3fc41e8
1 parent 4852e9d
commit 3fc41e8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/pytorch/pytorch_checkpoint.py b/pytorch/pytorch_checkpoint.py
@@ -104,6 +104,7 @@ def objective(trial):
             artifact_id=artifact_id,
         )
         checkpoint = torch.load(f"./tmp_model_{trial.number}.pt")
+        os.remove(f"./tmp_model_{trial.number}.pt")
         epoch = checkpoint["epoch"]
         epoch_begin = epoch + 1