allow errors during checkpoint loading (and

molmod · Jul 1, 2024 · 290e0b4 · 290e0b4
1 parent f70dbe2
commit 290e0b4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/psiflow/models/mace_utils.py b/psiflow/models/mace_utils.py
@@ -752,11 +752,16 @@ def run(rank: int, args: argparse.Namespace, world_size: int) -> None:
     }
 
     for swa_eval in swas:
-        epoch = checkpoint_handler.load_latest(
-            state=tools.CheckpointState(model, optimizer, lr_scheduler),
-            swa=swa_eval,
-            device=device,
-        )
+        try:
+            epoch = checkpoint_handler.load_latest(
+                state=tools.CheckpointState(model, optimizer, lr_scheduler),
+                swa=swa_eval,
+                device=device,
+            )
+        except BaseException as e:
+            print('failed to load checkpoint for swa:{}'.format(swa_eval))
+            print(e)
+            continue
         model.to(device)
         if args.distributed:
             distributed_model = DDP(model, device_ids=[local_rank])

diff --git a/tests/test_models.py b/tests/test_models.py
@@ -83,7 +83,7 @@ def test_mace_train(gpu, mace_config, dataset, tmp_path):
     # it with the manually computed value
     training = dataset[:-5]
     validation = dataset[-5:]
-    mace_config["start_swa"] = 1000
+    mace_config["start_swa"] = 100
     model = MACE(**mace_config)
     model.initialize(training)
     hamiltonian0 = model.create_hamiltonian()