From 784d7b9f3c45733c3b2081a4676aef2349ef3ca0 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 15 Mar 2024 11:30:26 +0800 Subject: [PATCH] fix: address comments --- deepmd/pt/model/model/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 9dd619ce34..7a2070e476 100644 --- 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"]