add docstring

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Dec 8, 2023 · ab5b3eb · ab5b3eb
1 parent 160929c
commit ab5b3eb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py
@@ -45,6 +45,9 @@ class DeepEval:
         as the initial batch size.
     input_map : dict, optional
         The input map for tf.import_graph_def. Only work with default tf graph
+    neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional
+        The ASE neighbor list class to produce the neighbor list. If None, the
+        neighbor list will be built natively in the model.
     """
 
     load_prefix: str  # set by subclass

diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py
@@ -51,6 +51,9 @@ class DeepPot(DeepEval):
         as the initial batch size.
     input_map : dict, optional
         The input map for tf.import_graph_def. Only work with default tf graph
+    neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional
+        The ASE neighbor list class to produce the neighbor list. If None, the
+        neighbor list will be built natively in the model.
 
     Examples
     --------