fix uts

deepmodeling · Jun 3, 2024 · ddaa38d · ddaa38d
1 parent aa08e30
commit ddaa38d
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 5 deletions.
diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py
@@ -183,6 +183,12 @@ def serialize(self) -> dict:
     def deserialize(cls) -> "EnergyFittingNetDirect":
         raise NotImplementedError
 
+    def slim_type_map(self, type_map: List[str]) -> None:
+        raise NotImplementedError
+
+    def get_type_map(self) -> List[str]:
+        raise NotImplementedError
+
     def forward(
         self,
         inputs: torch.Tensor,

diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py
@@ -657,7 +657,10 @@ def __init__(
             self.descrpt = descriptor
         else:
             self.descrpt = Descriptor(
-                **descriptor, ntypes=len(self.get_type_map()), spin=self.spin
+                **descriptor,
+                ntypes=len(self.get_type_map()),
+                spin=self.spin,
+                type_map=type_map,
             )
 
         if isinstance(fitting_net, Fitting):
@@ -672,6 +675,7 @@ def __init__(
                 ntypes=self.descrpt.get_ntypes(),
                 dim_descrpt=self.descrpt.get_dim_out(),
                 mixed_types=type_embedding is not None or self.descrpt.explicit_ntypes,
+                type_map=type_map,
             )
         self.rcut = self.descrpt.get_rcut()
         self.ntypes = self.descrpt.get_ntypes()
@@ -680,12 +684,11 @@ def __init__(
         if type_embedding is not None and isinstance(type_embedding, TypeEmbedNet):
             self.typeebd = type_embedding
         elif type_embedding is not None:
-            if type_embedding.get("use_econf_tebd", False):
-                type_embedding["type_map"] = type_map
             self.typeebd = TypeEmbedNet(
                 ntypes=self.ntypes,
                 **type_embedding,
                 padding=self.descrpt.explicit_ntypes,
+                type_map=type_map,
             )
         elif self.descrpt.explicit_ntypes:
             default_args = type_embedding_args()
@@ -695,6 +698,7 @@ def __init__(
                 ntypes=self.ntypes,
                 **default_args_dict,
                 padding=True,
+                type_map=type_map,
             )
         else:
             self.typeebd = None

diff --git a/deepmd/tf/model/pairwise_dprc.py b/deepmd/tf/model/pairwise_dprc.py
@@ -88,13 +88,12 @@ def __init__(
         if isinstance(type_embedding, TypeEmbedNet):
             self.typeebd = type_embedding
         else:
-            if type_embedding.get("use_econf_tebd", False):
-                type_embedding["type_map"] = type_map
             self.typeebd = TypeEmbedNet(
                 ntypes=self.ntypes,
                 **type_embedding,
                 # must use se_atten, so it must be True
                 padding=True,
+                type_map=type_map,
             )
 
         self.qm_model = Model(

diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py
@@ -112,6 +112,12 @@ def output_def(self):
     def serialize(self) -> dict:
         raise NotImplementedError
 
+    def slim_type_map(self, type_map: List[str]) -> None:
+        raise NotImplementedError
+
+    def get_type_map(self) -> List[str]:
+        raise NotImplementedError
+
     def forward(
         self,
         descriptor: torch.Tensor,