Update run.py

deepmodeling · May 9, 2024 · 083ad45 · 083ad45
1 parent 47a766e
commit 083ad45
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py
@@ -748,8 +748,8 @@ def run_train(iter_index, jdata, mdata):
             "training_init_model, training_init_frozen_model, and training_finetune_model are mutually exclusive."
         )
 
-    train_command = mdata.get("train_command", "dp")
-    assert train_command.strip() == "dp", "The 'train_command' should be 'dp'"
+    train_command = mdata.get("train_command", "dp").strip()
+    assert train_command == "dp", "The 'train_command' should be 'dp'"
     if suffix == ".pb":
         train_command += " --tf"
     elif suffix == ".pth":