Update multi_task.py

deepmodeling · Mar 5, 2024 · 37c6330 · 37c6330
1 parent 8376830
commit 37c6330
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/deepmd/pt/utils/multi_task.py b/deepmd/pt/utils/multi_task.py
@@ -143,8 +143,12 @@ def replace_one_item(params_dict, key_type, key_in_dict, suffix="", index=None):
                             )
     for shared_key in shared_links:
         shared_links[shared_key]["links"] = sorted(
-            shared_links[shared_key]["links"], key=lambda x: x["shared_level"]
+            shared_links[shared_key]["links"],
+            key=lambda x: x["shared_level"]
+            - ("spin" in model_config["model_dict"][x["model_key"]]) * 100,
         )
+        # little trick to make spin models in the front to be the base models,
+        # because its type embeddings are more general.
     assert len(type_map_keys) == 1, "Multitask model must have only one type_map!"
     return model_config, shared_links