model_type overwrite in config moved to right location in the code

ORNL · Nov 27, 2024 · 98d5a94 · 98d5a94
1 parent 52d8d5a
commit 98d5a94
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/examples/qm9/qm9.py b/examples/qm9/qm9.py
@@ -45,6 +45,10 @@ def main(model_type=None):
     with open(filename, "r") as f:
         config = json.load(f)
 
+    # If a model type is provided, update the configuration accordingly.
+    if model_type:
+        config["NeuralNetwork"]["Architecture"]["model_type"] = model_type
+
     verbosity = config["Verbosity"]["level"]
     var_config = config["NeuralNetwork"]["Variables_of_interest"]
 
@@ -73,10 +77,6 @@ def main(model_type=None):
         config, train_loader, val_loader, test_loader
     )
 
-    # If a model type is provided, update the configuration accordingly.
-    if model_type:
-        config["NeuralNetwork"]["Architecture"]["model_type"] = model_type
-
     model = hydragnn.models.create_model_config(
         config=config["NeuralNetwork"],
         verbosity=verbosity,