From 624d5928c4a02a2694de3a29388f8fd603f936eb Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Tue, 10 Jan 2023 17:56:11 -0500
Subject: [PATCH] fix dp freeze for multi-task when `atom_ener` is set

---
 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",