update model logic

nnaakkaaii · Jul 8, 2024 · 1704b1e · 1704b1e
1 parent 1a2a778
commit 1704b1e
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 7 deletions.
diff --git a/hrdae/conf b/hrdae/conf
diff --git a/hrdae/models/basic_model.py b/hrdae/models/basic_model.py
@@ -181,9 +181,6 @@ def train(
                     f"epoch_{epoch}",
                 )
 
-        with open(result_dir / "training_history.json", "w") as f:
-            json.dump(training_history, f)
-
         return least_val_loss
 
 

diff --git a/hrdae/models/vr_model.py b/hrdae/models/vr_model.py
@@ -144,6 +144,8 @@ def train(
                     "val_loss": float(avg_val_loss),
                 }
             )
+            with open(result_dir / "training_history.json", "w") as f:
+                json.dump(training_history, f)
 
             if epoch % 10 == 0:
                 data = next(iter(val_loader))
@@ -166,9 +168,6 @@ def train(
                     result_dir / "weights" / f"model_{epoch}.pth",
                 )
 
-        with open(result_dir / "training_history.json", "w") as f:
-            json.dump(training_history, f)
-
         return least_val_loss