fix non prediction bug

deepmodeling · Jan 5, 2025 · af1032a · af1032a
1 parent cb1d865
commit af1032a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/dptb/entrypoints/train.py b/dptb/entrypoints/train.py
@@ -192,7 +192,8 @@ def train(
         # build model will handle the init model cases where the model options provided is not equals to the ones in checkpoint.
         checkpoint = init_model if init_model else None
         model = build_model(checkpoint=checkpoint, model_options=jdata["model_options"], common_options=jdata["common_options"])
-        train_datasets.E3statistics(model=model, decay=jdata["model_options"]["prediction"]["decay"])
+        decay = jdata["model_options"].get("prediction", {}).get("decay", False)
+        train_datasets.E3statistics(model=model, decay=decay)
         trainer = Trainer(
             train_options=jdata["train_options"],
             common_options=jdata["common_options"],