Save current model.

usnistgov · Oct 2, 2023 · 60cd9bc · 60cd9bc
1 parent aa9c36b
commit 60cd9bc
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py
@@ -142,8 +142,8 @@ def get_figshare_model_ff(model_name="alignnff_fmult", dir_path=None):
 
 def default_path():
     """Get default model path."""
-    # dpath = get_figshare_model_ff(model_name="alignnff_wt10")
-    dpath = get_figshare_model_ff(model_name="alignnff_fmult")
+    dpath = get_figshare_model_ff(model_name="alignnff_wt10")
+    # dpath = get_figshare_model_ff(model_name="alignnff_fmult")
     print("model_path", dpath)
     return dpath
 

diff --git a/alignn/train.py b/alignn/train.py
@@ -647,6 +647,11 @@ def get_batch_errors(dat=[]):
             mean_out, mean_atom, mean_grad, mean_stress = get_batch_errors(
                 val_result
             )
+            current_model_name = "current_model.pt"
+            torch.save(
+                net.state_dict(),
+                os.path.join(config.output_dir, current_model_name),
+            )
             if val_loss < best_loss:
                 best_loss = val_loss
                 best_model_name = "best_model.pt"