fix dp freeze for multi-task when atom_ener is set (#2239)

When `atom_ener` is set, there will be variables named `layer_xx_suffix_1/MatMul` (the automatic name when `layer_xx_suffix/MatMul` has been used).
deepmodeling · Jan 11, 2023 · 30920ab · 30920ab
1 parent 6322609
commit 30920ab
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/deepmd/env.py b/deepmd/env.py
@@ -166,6 +166,8 @@ def dlopen_library(module: str, filename: str):
     "model_attr/sel_type_{}": "model_attr/sel_type",
     "model_attr/output_dim_{}": "model_attr/output_dim",
     "_{}/": "/",
+    # when atom_ener is set
+    "_{}_1/": "_1/",
     "o_energy_{}": "o_energy",
     "o_force_{}": "o_force",
     "o_virial_{}": "o_virial",