Save a checkpoint at the end of every epoch.

bd34e4d4 · Wes Brown · 55a00bbf · bd34e4d4
Commit bd34e4d4 authored Jul 18, 2022 by Wes Brown
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

hypertrain.py hypertrain.py +2 -0

No files found.
--- a/hypertrain.py
+++ b/hypertrain.py
@@ -444,6 +444,8 @@ with tqdm(total=total_steps, initial=curr_step) as t:

            curr_step += 1
            t.update(1)
+        if train_config["epochs"] > 1:
+            hypernetwork_saver(f"epoch-{epoch}")

 eval_fn(curr_step)
 hypernetwork_saver("final")