fix docstr

Author: iProzd

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -363,10 +363,10 @@ def compute_or_load_stat(
         self,
         merged: Union[Callable[[], list[dict]], list[dict]],
         stat_file_path: Optional[DPPath] = None,
-        compute_out_stat: bool = True,
+        compute_or_load_out_stat: bool = True,
     ) -> NoReturn:
         """
-        Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
+        Compute the input and output statistics (e.g. energy bias) for the model from packed data.
 
         Parameters
         ----------
@@ -379,6 +379,9 @@ def compute_or_load_stat(
                 the lazy function helps by only sampling once.
         stat_file_path : Optional[DPPath]
             The path to the stat file.
+        compute_or_load_out_stat : bool
+            Whether to compute the output statistics.
+            If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors).
 
         """
         raise NotImplementedError

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -285,7 +285,7 @@ def compute_or_load_stat(
         self,
         sampled_func,
         stat_file_path: Optional[DPPath] = None,
-        compute_out_stat: bool = True,
+        compute_or_load_out_stat: bool = True,
     ) -> None:
         """
         Compute or load the statistics parameters of the model,
@@ -301,6 +301,9 @@ def compute_or_load_stat(
             The lazy sampled function to get data frames from different data systems.
         stat_file_path
             The dictionary of paths to the statistics files.
+        compute_or_load_out_stat : bool
+            Whether to compute the output statistics.
+            If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors).
         """
         if stat_file_path is not None and self.type_map is not None:
             # descriptors and fitting net with different type_map
@@ -324,7 +327,7 @@ def wrapped_sampler():
         self.fitting_net.compute_input_stats(
             wrapped_sampler, protection=self.data_stat_protect
         )
-        if compute_out_stat:
+        if compute_or_load_out_stat:
             self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
     def get_dim_fparam(self) -> int:

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -472,7 +472,7 @@ def compute_or_load_stat(
         self,
         sampled_func,
         stat_file_path: Optional[DPPath] = None,
-        compute_out_stat: bool = True,
+        compute_or_load_out_stat: bool = True,
     ) -> None:
         """
         Compute or load the statistics parameters of the model,
@@ -488,10 +488,13 @@ def compute_or_load_stat(
             The lazy sampled function to get data frames from different data systems.
         stat_file_path
             The dictionary of paths to the statistics files.
+        compute_or_load_out_stat : bool
+            Whether to compute the output statistics.
+            If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors).
         """
         for md in self.models:
             md.compute_or_load_stat(
-                sampled_func, stat_file_path, compute_out_stat=False
+                sampled_func, stat_file_path, compute_or_load_out_stat=False
             )
 
         if stat_file_path is not None and self.type_map is not None:

deepmd/pt/model/atomic_model/pairtab_atomic_model.py

@@ -226,10 +226,10 @@ def compute_or_load_stat(
         self,
         merged: Union[Callable[[], list[dict]], list[dict]],
         stat_file_path: Optional[DPPath] = None,
-        compute_out_stat: bool = True,
+        compute_or_load_out_stat: bool = True,
     ) -> None:
         """
-        Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
+        Compute the input and output statistics (e.g. energy bias) for the model from packed data.
 
         Parameters
         ----------
@@ -242,9 +242,12 @@ def compute_or_load_stat(
                 the lazy function helps by only sampling once.
         stat_file_path : Optional[DPPath]
             The path to the stat file.
+        compute_or_load_out_stat : bool
+            Whether to compute the output statistics.
+            If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors).
 
         """
-        if compute_out_stat:
+        if compute_or_load_out_stat:
             self.compute_or_load_out_stat(merged, stat_file_path)
 
     def forward_atomic(