Update deepmd/dpmodel/utils/type_embed.py

Co-authored-by: Han Wang <[email protected]> Signed-off-by: Duo <[email protected]>
deepmodeling · May 15, 2024 · d2e0db0 · d2e0db0
1 parent 0cc6af9
commit d2e0db0
Showing 1 changed file with 7 additions and 15 deletions.
diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py
@@ -87,21 +87,13 @@ def __init__(
                 [electronic_configuration_embedding[kk] for kk in self.type_map],
                 dtype=PRECISION_DICT[self.precision],
             )
-            self.embedding_net = EmbeddingNet(
-                ECONF_DIM,
-                self.neuron,
-                self.activation_function,
-                self.resnet_dt,
-                self.precision,
-            )
-        else:
-            self.embedding_net = EmbeddingNet(
-                ntypes,
-                self.neuron,
-                self.activation_function,
-                self.resnet_dt,
-                self.precision,
-            )
+        self.embedding_net = EmbeddingNet(
+            ECONF_DIM if self.use_econf_tebd else ntypes,
+            self.neuron,
+            self.activation_function,
+            self.resnet_dt,
+            self.precision,
+        )
 
     def call(self) -> np.ndarray:
         """Compute the type embedding network."""