From 30920abe0e843e1fc04b5fb02af17cf0b4ff31f2 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Tue, 10 Jan 2023 23:06:27 -0500
Subject: [PATCH] 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).
---
 deepmd/env.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/deepmd/env.py b/deepmd/env.py
index 74809be39b..be074bb162 100644
--- 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",