add missing docs in pt make model

wanghan-iapcm · Feb 11, 2024 · cb7b96d · cb7b96d
1 parent c57dccc
commit cb7b96d
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
@@ -83,6 +83,10 @@ def forward_common(
                 The type of atoms. shape: nf x nloc
             box
                 The simulation box. shape: nf x 9
+            fparam
+                frame parameter. nf x ndf
+            aparam
+                atomic parameter. nf x nloc x nda
             do_atomic_virial
                 If calculate the atomic virial.
 
@@ -155,6 +159,10 @@ def forward_common_lower(
                 neighbor list. nf x nloc x nsel.
             mapping
                 mapps the extended indices to local indices. nf x nall.
+            fparam
+                frame parameter. nf x ndf
+            aparam
+                atomic parameter. nf x nloc x nda
             do_atomic_virial
                 whether calculate atomic virial.