fix the bug of empty fitting net neuron. (deepmodeling#3458)

fixes deepmodeling#3448 Co-authored-by: Han Wang <[email protected]>
njzjz · Mar 13, 2024 · 571ddec · 571ddec
1 parent da68686
commit 571ddec
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py
@@ -609,7 +609,8 @@ def __init__(
                 resnet_dt=resnet_dt,
                 precision=precision,
             )
-            i_in, i_ot = neuron[-1], out_dim
+            i_in = neuron[-1] if len(neuron) > 0 else in_dim
+            i_ot = out_dim
             self.layers.append(
                 T_NetworkLayer(
                     i_in,

diff --git a/source/tests/pt/model/test_ener_fitting.py b/source/tests/pt/model/test_ener_fitting.py
@@ -44,12 +44,13 @@ def test_consistency(
         )
         atype = torch.tensor(self.atype_ext[:, :nloc], dtype=int, device=env.DEVICE)
 
-        for od, mixed_types, nfp, nap, et in itertools.product(
+        for od, mixed_types, nfp, nap, et, nn in itertools.product(
             [1, 3],
             [True, False],
             [0, 3],
             [0, 4],
             [[], [0], [1]],
+            [[4, 4, 4], []],
         ):
             ft0 = InvarFitting(
                 "foo",
@@ -60,6 +61,7 @@ def test_consistency(
                 numb_aparam=nap,
                 mixed_types=mixed_types,
                 exclude_types=et,
+                neuron=nn,
             ).to(env.DEVICE)
             ft1 = DPInvarFitting.deserialize(ft0.serialize())
             ft2 = InvarFitting.deserialize(ft0.serialize())