diff --git a/source/tests/common/dpmodel/test_descriptor_dpa2.py b/source/tests/common/dpmodel/test_descriptor_dpa2.py
index 3ae0689dad..cf6789faf6 100644
--- a/source/tests/common/dpmodel/test_descriptor_dpa2.py
+++ b/source/tests/common/dpmodel/test_descriptor_dpa2.py
@@ -6,6 +6,10 @@
 from deepmd.dpmodel.descriptor import (
     DescrptDPA2,
 )
+from deepmd.dpmodel.descriptor.dpa2 import (
+    RepformerArgs,
+    RepinitArgs,
+)
 
 from .case_single_frame_with_nlist import (
     TestCaseSingleFrameWithNlist,
@@ -28,14 +32,21 @@ def test_self_consistency(
         dstd = 0.1 + np.abs(dstd)
         dstd_2 = 0.1 + np.abs(dstd_2)
 
+        repinit = RepinitArgs(
+            rcut=self.rcut,
+            rcut_smth=self.rcut_smth,
+            nsel=self.sel_mix,
+        )
+        repformer = RepformerArgs(
+            rcut=self.rcut / 2,
+            rcut_smth=self.rcut_smth,
+            nsel=nnei // 2,
+        )
+
         em0 = DescrptDPA2(
             ntypes=self.nt,
-            repinit_rcut=self.rcut,
-            repinit_rcut_smth=self.rcut_smth,
-            repinit_nsel=self.sel_mix,
-            repformer_rcut=self.rcut / 2,
-            repformer_rcut_smth=self.rcut_smth,
-            repformer_nsel=nnei // 2,
+            repinit=repinit,
+            repformer=repformer,
         )
 
         em0.repinit.mean = davg
diff --git a/source/tests/pt/model/models/dpa2.json b/source/tests/pt/model/models/dpa2.json
index 966e298dfa..ca1948492a 100644
--- a/source/tests/pt/model/models/dpa2.json
+++ b/source/tests/pt/model/models/dpa2.json
@@ -5,35 +5,40 @@
   ],
   "descriptor": {
     "type": "dpa2",
-    "repinit_rcut": 6.0,
-    "repinit_rcut_smth": 2.0,
-    "repinit_nsel": 30,
-    "repformer_rcut": 4.0,
-    "repformer_rcut_smth": 0.5,
-    "repformer_nsel": 10,
-    "repinit_neuron": [
-      2,
-      4,
-      8
-    ],
-    "repinit_axis_neuron": 4,
-    "repinit_activation_function": "tanh",
-    "repformer_nlayers": 12,
-    "repformer_g1_dim": 8,
-    "repformer_g2_dim": 5,
-    "repformer_attn2_hidden": 3,
-    "repformer_attn2_nhead": 1,
-    "repformer_attn1_hidden": 5,
-    "repformer_attn1_nhead": 1,
-    "repformer_axis_neuron": 4,
-    "repformer_update_h2": false,
-    "repformer_update_g1_has_conv": true,
-    "repformer_update_g1_has_grrg": true,
-    "repformer_update_g1_has_drrd": true,
-    "repformer_update_g1_has_attn": true,
-    "repformer_update_g2_has_g1g1": true,
-    "repformer_update_g2_has_attn": true,
-    "repformer_attn2_has_gate": true,
+    "repinit": {
+      "rcut": 6.0,
+      "rcut_smth": 2.0,
+      "nsel": 30,
+      "neuron": [
+        2,
+        4,
+        8
+      ],
+      "axis_neuron": 4,
+      "activation_function": "tanh"
+
+    },
+    "repformer": {
+      "rcut": 4.0,
+      "rcut_smth": 0.5,
+      "nsel": 10,
+      "nlayers": 12,
+      "g1_dim": 8,
+      "g2_dim": 5,
+      "attn2_hidden": 3,
+      "attn2_nhead": 1,
+      "attn1_hidden": 5,
+      "attn1_nhead": 1,
+      "axis_neuron": 4,
+      "update_h2": false,
+      "update_g1_has_conv": true,
+      "update_g1_has_grrg": true,
+      "update_g1_has_drrd": true,
+      "update_g1_has_attn": true,
+      "update_g2_has_g1g1": true,
+      "update_g2_has_attn": true,
+      "attn2_has_gate": true
+    },
     "add_tebd_to_repinit_out": false
   },
   "fitting_net": {