Fix uts

deepmodeling · Mar 1, 2024 · a953f7a · a953f7a
1 parent a8f58a7
commit a953f7a
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py
@@ -295,7 +295,7 @@ class EnergyFittingNetDirect(Fitting):
     def __init__(
         self,
         ntypes,
-        embedding_width,
+        dim_descrpt,
         neuron,
         bias_atom_e=None,
         out_dim=1,
@@ -315,7 +315,7 @@ def __init__(
         """
         super().__init__()
         self.ntypes = ntypes
-        self.dim_descrpt = embedding_width
+        self.dim_descrpt = dim_descrpt
         self.use_tebd = use_tebd
         self.out_dim = out_dim
         if bias_atom_e is None:
@@ -329,7 +329,7 @@ def __init__(
         for type_i in range(self.ntypes):
             one = ResidualDeep(
                 type_i,
-                embedding_width,
+                dim_descrpt,
                 neuron,
                 0.0,
                 out_dim=out_dim,
@@ -344,7 +344,7 @@ def __init__(
             for type_i in range(self.ntypes):
                 bias_type = 0.0 if self.use_tebd else bias_atom_e[type_i]
                 one = ResidualDeep(
-                    type_i, embedding_width, neuron, bias_type, resnet_dt=resnet_dt
+                    type_i, dim_descrpt, neuron, bias_type, resnet_dt=resnet_dt
                 )
                 filter_layers.append(one)
         self.filter_layers = torch.nn.ModuleList(filter_layers)

diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py
@@ -33,6 +33,8 @@
     p_examples / "nopbc" / "train" / "input.json",
     p_examples / "water_tensor" / "dipole" / "dipole_input.json",
     p_examples / "water_tensor" / "polar" / "polar_input.json",
+    p_examples / "water_tensor" / "dipole" / "dipole_input_torch.json",
+    p_examples / "water_tensor" / "polar" / "polar_input_torch.json",
     p_examples / "water_multi_task" / "ener_dipole" / "input.json",
     p_examples / "fparam" / "train" / "input.json",
     p_examples / "fparam" / "train" / "input_aparam.json",

diff --git a/source/tests/pt/water_tensor/se_e2_a.json b/source/tests/pt/water_tensor/se_e2_a.json
@@ -44,7 +44,8 @@
   },
   "learning_rate": {
     "type": "exp",
-    "start_lr": 0.01,
+    "start_lr": 0.001,
+    "stop_lr": 3.51e-8,
     "decay_steps": 5000,
     "_comment5": "that's all"
   },