diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py
index fec04255fa..f507374ff2 100644
--- a/deepmd/dpmodel/model/make_model.py
+++ b/deepmd/dpmodel/model/make_model.py
@@ -149,13 +149,13 @@ def call_lower(
             Parameters
             ----------
             extended_coord
-                coodinates in extended region
+                coodinates in extended region. nf x (nall x 3).
             extended_atype
-                atomic type in extended region
+                atomic type in extended region. nf x nall.
             nlist
-                neighbor list. nf x nloc x nsel
+                neighbor list. nf x nloc x nsel.
             mapping
-                mapps the extended indices to local indices
+                mapps the extended indices to local indices. nf x nall.
             fparam
                 frame parameter. nf x ndf
             aparam
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index c8c1e9450b..1e76c6a468 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -147,15 +147,15 @@ def forward_common_lower(
             Parameters
             ----------
             extended_coord
-                coodinates in extended region
+                coodinates in extended region. nf x (nall x 3)
             extended_atype
-                atomic type in extended region
+                atomic type in extended region. nf x nall
             nlist
-                neighbor list. nf x nloc x nsel
+                neighbor list. nf x nloc x nsel.
             mapping
-                mapps the extended indices to local indices
+                mapps the extended indices to local indices. nf x nall.
             do_atomic_virial
-                whether calculate atomic virial
+                whether calculate atomic virial.
 
             Returns
             -------