From ab5b3eb4be5c88281f0d1a17fd96e70fdeb7967d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Fri, 8 Dec 2023 03:17:08 -0500 Subject: [PATCH] add docstring Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_eval.py | 3 +++ deepmd/infer/deep_pot.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index d61227f2ef..0ca9f21a77 100644 --- 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 index 6a658b14fd..81cfdde7a8 100644 --- 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 --------