fix config file test to not require prior_model in the configs

torchmd · Feb 19, 2024 · 1bbb03a · 1bbb03a
1 parent e79b9c0
commit 1bbb03a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -19,7 +19,7 @@ def test_example_yamls(fname):
         args = yaml.load(f, Loader=yaml.FullLoader)
 
     prior = None
-    if args["prior_model"] is not None:
+    if "prior_model" in args and args["prior_model"] is not None:
         dataset = DummyDataset(has_atomref=True)
         prior = getattr(priors, args["prior_model"])(dataset=dataset)