fix ut

deepmodeling · May 12, 2024 · ada7f89 · ada7f89
1 parent 7936f9d
commit ada7f89
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py
@@ -691,8 +691,8 @@ def deserialize(cls, data: dict) -> "DescrptDPA2":
         g1_shape_tranform = data.pop("g1_shape_tranform")
         tebd_transform = data.pop("tebd_transform", None)
         add_tebd_to_repinit_out = data["add_tebd_to_repinit_out"]
-        data["repinit"] = RepinitArgs(**data["repinit_args"])
-        data["repformer"] = RepformerArgs(**data["repformer_args"])
+        data["repinit"] = RepinitArgs(**data.pop("repinit_args"))
+        data["repformer"] = RepformerArgs(**data.pop("repformer_args"))
         obj = cls(**data)
         obj.type_embedding = TypeEmbedNet.deserialize(type_embedding)
         if add_tebd_to_repinit_out:

diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py
@@ -418,8 +418,8 @@ def deserialize(cls, data: dict) -> "DescrptDPA2":
         g1_shape_tranform = data.pop("g1_shape_tranform")
         tebd_transform = data.pop("tebd_transform", None)
         add_tebd_to_repinit_out = data["add_tebd_to_repinit_out"]
-        data["repinit"] = RepinitArgs(**data["repinit_args"])
-        data["repformer"] = RepformerArgs(**data["repformer_args"])
+        data["repinit"] = RepinitArgs(**data.pop("repinit_args"))
+        data["repformer"] = RepformerArgs(**data.pop("repformer_args"))
         obj = cls(**data)
         obj.type_embedding.embedding = TypeEmbedNetConsistent.deserialize(
             type_embedding