add template_slice_config to run_train_config (#193)

Signed-off-by: zjgemi <liuxin_zijian@163.com>
deepmodeling · Feb 21, 2024 · 03fb112 · 03fb112
1 parent c0f1973
commit 03fb112
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/dpgen2/superop/prep_run_dp_train.py b/dpgen2/superop/prep_run_dp_train.py
@@ -232,6 +232,7 @@ def _prep_run_dp_train(
     run_template_config = run_config.pop("template_config")
     prep_executor = init_executor(prep_config.pop("executor"))
     run_executor = init_executor(run_config.pop("executor"))
+    template_slice_config = run_config.pop("template_slice_config", {})
 
     prep_train = Step(
         "prep-train",
@@ -261,6 +262,7 @@ def _prep_run_dp_train(
                 input_parameter=["task_name"],
                 input_artifact=["task_path", "init_model"],
                 output_artifact=["model", "lcurve", "log", "script"],
+                **template_slice_config,
             ),
             python_packages=upload_python_packages,
             **run_template_config,