fix: address comments

deepmodeling · Mar 15, 2024 · 784d7b9 · 784d7b9
1 parent 60315ee
commit 784d7b9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py
@@ -140,7 +140,7 @@ def get_standard_model(model_params):
     fitting_net["type"] = fitting_net.get("type", "ener")
     fitting_net["ntypes"] = descriptor.get_ntypes()
     fitting_net["mixed_types"] = descriptor.mixed_types()
-    if fitting_net["type"] != "dos":
+    if fitting_net["type"] in ["dipole", "polar"]:
         fitting_net["embedding_width"] = descriptor.get_dim_emb()
     fitting_net["dim_descrpt"] = descriptor.get_dim_out()
     grad_force = "direct" not in fitting_net["type"]