fix tests

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Feb 17, 2024 · 3bcdc33 · 3bcdc33
1 parent bca0f44
commit 3bcdc33
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py
@@ -407,6 +407,7 @@ def __init__(
         use_aparam_as_mask: bool = False,
         spin: Any = None,
         distinguish_types: bool = False,
+        exclude_types: List[int] = [],
         # not used
         seed: Optional[int] = None,
     ):

diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py
@@ -985,6 +985,7 @@ def serialize(self, suffix: str) -> dict:
             "layer_name": self.layer_name,
             "use_aparam_as_mask": self.use_aparam_as_mask,
             "spin": self.spin,
+            "exclude_types": [],
             "nets": self.serialize_network(
                 ntypes=self.ntypes,
                 # TODO: consider type embeddings

diff --git a/source/tests/consistent/fitting/test_ener.py b/source/tests/consistent/fitting/test_ener.py
@@ -150,7 +150,7 @@ def eval_pt(self, pt_obj: Any) -> Any:
         return (
             pt_obj(
                 torch.from_numpy(self.inputs).to(device=PT_DEVICE),
-                torch.from_numpy(self.atype).to(device=PT_DEVICE),
+                torch.from_numpy(self.atype.reshape(1, -1)).to(device=PT_DEVICE),
                 fparam=torch.from_numpy(self.fparam).to(device=PT_DEVICE)
                 if numb_fparam
                 else None,
@@ -169,7 +169,7 @@ def eval_dp(self, dp_obj: Any) -> Any:
         ) = self.param
         return dp_obj(
             self.inputs,
-            self.atype,
+            self.atype.reshape(1, -1),
             fparam=self.fparam if numb_fparam else None,
         )["energy"]