diff --git a/deepmd/pt/utils/multi_task.py b/deepmd/pt/utils/multi_task.py
index 5f06d93208..e2076b3b2b 100644
--- a/deepmd/pt/utils/multi_task.py
+++ b/deepmd/pt/utils/multi_task.py
@@ -115,7 +115,7 @@ def replace_one_item(params_dict, key_type, key_in_dict, suffix="", index=None):
                 type_map_keys.append(key_in_dict)
         else:
             if shared_key not in shared_links:
-                class_name = get_class_name(shared_type, shared_dict[key_in_dict])
+                class_name = get_class_name(shared_type, shared_dict[shared_key])
                 shared_links[shared_key] = {"type": class_name, "links": []}
             link_item = {
                 "model_key": model_key,