diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py
index 138cab82d6..42929ca1d3 100644
--- a/deepmd/dpmodel/descriptor/dpa1.py
+++ b/deepmd/dpmodel/descriptor/dpa1.py
@@ -14,6 +14,9 @@
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -524,18 +527,36 @@ def deserialize(cls, data: dict) -> "DescrptDPA1":
         return obj
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True)
+        min_nbor_dist, sel = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True
+        )
+        local_jdata_cpy["sel"] = sel[0]
+        return local_jdata_cpy, min_nbor_dist
 
 
 @DescriptorBlock.register("se_atten")
diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py
index 78bf174685..9243883ebe 100644
--- a/deepmd/dpmodel/descriptor/dpa2.py
+++ b/deepmd/dpmodel/descriptor/dpa2.py
@@ -15,6 +15,9 @@
 from deepmd.dpmodel.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -743,30 +746,46 @@ def deserialize(cls, data: dict) -> "DescrptDPA2":
         return obj
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
         update_sel = UpdateSel()
-        local_jdata_cpy["repinit"] = update_sel.update_one_sel(
-            global_jdata,
-            local_jdata_cpy["repinit"],
+        min_nbor_dist, repinit_sel = update_sel.update_one_sel(
+            train_data,
+            type_map,
+            local_jdata_cpy["repinit"]["rcut"],
+            local_jdata_cpy["repinit"]["nsel"],
             True,
-            rcut_key="rcut",
-            sel_key="nsel",
         )
-        local_jdata_cpy["repformer"] = update_sel.update_one_sel(
-            global_jdata,
-            local_jdata_cpy["repformer"],
+        local_jdata_cpy["repinit"]["nsel"] = repinit_sel[0]
+        min_nbor_dist, repformer_sel = update_sel.update_one_sel(
+            train_data,
+            type_map,
+            local_jdata_cpy["repformer"]["rcut"],
+            local_jdata_cpy["repformer"]["nsel"],
             True,
-            rcut_key="rcut",
-            sel_key="nsel",
         )
-        return local_jdata_cpy
+        local_jdata_cpy["repformer"]["nsel"] = repformer_sel[0]
+        return local_jdata_cpy, min_nbor_dist
diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py
index 15825ecc10..d359bf6fb7 100644
--- a/deepmd/dpmodel/descriptor/hybrid.py
+++ b/deepmd/dpmodel/descriptor/hybrid.py
@@ -5,6 +5,7 @@
     Dict,
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -19,6 +20,9 @@
 from deepmd.dpmodel.utils.nlist import (
     nlist_distinguish_types,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -222,22 +226,42 @@ def call(
         return out_descriptor, out_gr, out_g2, out_h2, out_sw
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict:
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        local_jdata_cpy["list"] = [
-            BaseDescriptor.update_sel(global_jdata, sub_jdata)
-            for sub_jdata in local_jdata["list"]
-        ]
-        return local_jdata_cpy
+        new_list = []
+        min_nbor_dist = None
+        for sub_jdata in local_jdata["list"]:
+            new_sub_jdata, min_nbor_dist_ = BaseDescriptor.update_sel(
+                train_data, type_map, sub_jdata
+            )
+            if min_nbor_dist_ is not None:
+                min_nbor_dist = min_nbor_dist_
+            new_list.append(new_sub_jdata)
+        local_jdata_cpy["list"] = new_list
+        return local_jdata_cpy, min_nbor_dist
 
     def serialize(self) -> dict:
         return {
diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py
index cba9eebe4b..a4fc8bddf9 100644
--- a/deepmd/dpmodel/descriptor/make_base_descriptor.py
+++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py
@@ -7,12 +7,16 @@
     Callable,
     List,
     Optional,
+    Tuple,
     Union,
 )
 
 from deepmd.common import (
     j_get_type,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -149,19 +153,33 @@ def deserialize(cls, data: dict) -> "BD":
 
         @classmethod
         @abstractmethod
-        def update_sel(cls, global_jdata: dict, local_jdata: dict):
+        def update_sel(
+            cls,
+            train_data: DeepmdDataSystem,
+            type_map: Optional[List[str]],
+            local_jdata: dict,
+        ) -> Tuple[dict, Optional[float]]:
             """Update the selection and perform neighbor statistics.
 
             Parameters
             ----------
-            global_jdata : dict
-                The global data, containing the training section
+            train_data : DeepmdDataSystem
+                data used to do neighbor statictics
+            type_map : list[str], optional
+                The name of each type of atoms
             local_jdata : dict
                 The local data refer to the current class
+
+            Returns
+            -------
+            dict
+                The updated local data
+            float
+                The minimum distance between two atoms
             """
             # call subprocess
             cls = cls.get_class_by_type(j_get_type(local_jdata, cls.__name__))
-            return cls.update_sel(global_jdata, local_jdata)
+            return cls.update_sel(train_data, type_map, local_jdata)
 
     setattr(BD, fwd_method_name, BD.fwd)
     delattr(BD, "fwd")
diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py
index 193383ac4f..2b6c3a843e 100644
--- a/deepmd/dpmodel/descriptor/se_e2_a.py
+++ b/deepmd/dpmodel/descriptor/se_e2_a.py
@@ -9,6 +9,9 @@
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -422,15 +425,32 @@ def deserialize(cls, data: dict) -> "DescrptSeA":
         return obj
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False)
+        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
+        )
+        return local_jdata_cpy, min_nbor_dist
diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py
index 5175b91ae1..e4a0a80657 100644
--- a/deepmd/dpmodel/descriptor/se_r.py
+++ b/deepmd/dpmodel/descriptor/se_r.py
@@ -4,6 +4,9 @@
 from deepmd.dpmodel.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -21,6 +24,7 @@
     Any,
     List,
     Optional,
+    Tuple,
 )
 
 from deepmd.dpmodel import (
@@ -345,15 +349,32 @@ def deserialize(cls, data: dict) -> "DescrptSeR":
         return obj
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False)
+        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
+        )
+        return local_jdata_cpy, min_nbor_dist
diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py
index d9b2741b3e..eac6a9640e 100644
--- a/deepmd/dpmodel/descriptor/se_t.py
+++ b/deepmd/dpmodel/descriptor/se_t.py
@@ -9,6 +9,9 @@
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -348,15 +351,32 @@ def deserialize(cls, data: dict) -> "DescrptSeT":
         return obj
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False)
+        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
+        )
+        return local_jdata_cpy, min_nbor_dist
diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py
index b416e57fb8..354fa58348 100644
--- a/deepmd/dpmodel/model/base_model.py
+++ b/deepmd/dpmodel/model/base_model.py
@@ -7,9 +7,14 @@
 from typing import (
     Any,
     List,
+    Optional,
+    Tuple,
     Type,
 )
 
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.plugin import (
     PluginVariant,
     make_plugin_registry,
@@ -129,12 +134,21 @@ def deserialize(cls, data: dict) -> "BaseBaseModel":
             raise NotImplementedError(f"Not implemented in class {cls.__name__}")
 
         model_def_script: str
+        """The model definition script."""
+        min_nbor_dist: Optional[float]
+        """The minimum distance between two atoms. Used for model compression.
+        None when skipping neighbor statistics.
+        """
 
         @abstractmethod
         def get_model_def_script(self) -> str:
             """Get the model definition script."""
             pass
 
+        def get_min_nbor_dist(self) -> Optional[float]:
+            """Get the minimum distance between two atoms."""
+            return self.min_nbor_dist
+
         @abstractmethod
         def get_nnei(self) -> int:
             """Returns the total number of selected neighboring atoms in the cut-off radius."""
@@ -148,22 +162,36 @@ def get_nsel(self) -> int:
 
         @classmethod
         @abstractmethod
-        def update_sel(cls, global_jdata: dict, local_jdata: dict):
+        def update_sel(
+            cls,
+            train_data: DeepmdDataSystem,
+            type_map: Optional[List[str]],
+            local_jdata: dict,
+        ) -> Tuple[dict, Optional[float]]:
             """Update the selection and perform neighbor statistics.
 
             Parameters
             ----------
-            global_jdata : dict
-                The global data, containing the training section
+            train_data : DeepmdDataSystem
+                data used to do neighbor statictics
+            type_map : list[str], optional
+                The name of each type of atoms
             local_jdata : dict
                 The local data refer to the current class
+
+            Returns
+            -------
+            dict
+                The updated local data
+            float
+                The minimum distance between two atoms
             """
             # getting model type based on fitting type
             model_type = local_jdata.get("type", "standard")
             if model_type == "standard":
                 model_type = local_jdata.get("fitting", {}).get("type", "ener")
             cls = cls.get_class_by_type(model_type)
-            return cls.update_sel(global_jdata, local_jdata)
+            return cls.update_sel(train_data, type_map, local_jdata)
 
     return BaseBaseModel
 
@@ -186,6 +214,7 @@ class BaseModel(make_base_model()):
 
     def __init__(self) -> None:
         self.model_def_script = ""
+        self.min_nbor_dist = None
 
     def get_model_def_script(self) -> str:
         """Get the model definition script."""
diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py
index 37cb426ab7..1597ba0b14 100644
--- a/deepmd/dpmodel/model/dp_model.py
+++ b/deepmd/dpmodel/model/dp_model.py
@@ -1,26 +1,49 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
 
+from typing import (
+    List,
+    Optional,
+    Tuple,
+)
+
 from deepmd.dpmodel.descriptor.base_descriptor import (
     BaseDescriptor,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 
 # use "class" to resolve "Variable not allowed in type expression"
 class DPModelCommon:
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel(
-            global_jdata, local_jdata["descriptor"]
+        local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel(
+            train_data, type_map, local_jdata["descriptor"]
         )
-        return local_jdata_cpy
+        return local_jdata_cpy, min_nbor_dist
diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py
index 90e2bb3fb4..c2cea35d27 100644
--- a/deepmd/dpmodel/model/spin_model.py
+++ b/deepmd/dpmodel/model/spin_model.py
@@ -227,6 +227,10 @@ def get_model_def_script(self) -> str:
         """Get the model definition script."""
         return self.backbone_model.get_model_def_script()
 
+    def get_min_nbor_dist(self) -> Optional[float]:
+        """Get the minimum neighbor distance."""
+        return self.backbone_model.get_min_nbor_dist()
+
     def get_nnei(self) -> int:
         """Returns the total number of selected neighboring atoms in the cut-off radius."""
         # for C++ interface
diff --git a/deepmd/dpmodel/utils/update_sel.py b/deepmd/dpmodel/utils/update_sel.py
index 48463b5743..dc38a6a041 100644
--- a/deepmd/dpmodel/utils/update_sel.py
+++ b/deepmd/dpmodel/utils/update_sel.py
@@ -15,7 +15,3 @@ class UpdateSel(BaseUpdateSel):
     @property
     def neighbor_stat(self) -> Type[NeighborStat]:
         return NeighborStat
-
-    def hook(self, min_nbor_dist, max_nbor_size):
-        # TODO: save to the model in UpdateSel.hook
-        pass
diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py
index fba22e6d24..991278f0aa 100644
--- a/deepmd/pt/entrypoints/main.py
+++ b/deepmd/pt/entrypoints/main.py
@@ -3,9 +3,6 @@
 import json
 import logging
 import os
-from copy import (
-    deepcopy,
-)
 from pathlib import (
     Path,
 )
@@ -66,6 +63,7 @@
     update_deepmd_input,
 )
 from deepmd.utils.data_system import (
+    get_data,
     process_systems,
 )
 from deepmd.utils.path import (
@@ -251,24 +249,33 @@ def train(FLAGS):
         config = normalize(config)
 
     # do neighbor stat
+    min_nbor_dist = None
     if not FLAGS.skip_neighbor_stat:
         log.info(
             "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)"
         )
+
+        type_map = config["model"].get("type_map")
         if not multi_task:
-            config["model"] = BaseModel.update_sel(config, config["model"])
+            train_data = get_data(
+                config["training"]["training_data"], 0, type_map, None
+            )
+            config["model"], min_nbor_dist = BaseModel.update_sel(
+                train_data, type_map, config["model"]
+            )
         else:
-            training_jdata = deepcopy(config["training"])
-            training_jdata.pop("data_dict", {})
-            training_jdata.pop("model_prob", {})
+            min_nbor_dist = {}
             for model_item in config["model"]["model_dict"]:
-                fake_global_jdata = {
-                    "model": deepcopy(config["model"]["model_dict"][model_item]),
-                    "training": deepcopy(config["training"]["data_dict"][model_item]),
-                }
-                fake_global_jdata["training"].update(training_jdata)
-                config["model"]["model_dict"][model_item] = BaseModel.update_sel(
-                    fake_global_jdata, config["model"]["model_dict"][model_item]
+                train_data = get_data(
+                    config["training"]["data_dict"][model_item]["training_data"],
+                    0,
+                    type_map,
+                    None,
+                )
+                config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = (
+                    BaseModel.update_sel(
+                        train_data, type_map, config["model"]["model_dict"][model_item]
+                    )
                 )
 
     with open(FLAGS.output, "w") as fp:
@@ -284,6 +291,13 @@ def train(FLAGS):
         FLAGS.init_frz_model,
         shared_links=shared_links,
     )
+    # save min_nbor_dist
+    if min_nbor_dist is not None:
+        if not multi_task:
+            trainer.model.min_nbor_dist = min_nbor_dist
+        else:
+            for model_item in min_nbor_dist:
+                trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item]
     trainer.run()
 
 
diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py
index 4ab39465dc..92f5cf2e15 100644
--- a/deepmd/pt/model/descriptor/dpa1.py
+++ b/deepmd/pt/model/descriptor/dpa1.py
@@ -27,6 +27,9 @@
 from deepmd.pt.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -549,15 +552,33 @@ def forward(
         return g1, rot_mat, g2, h2, sw
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True)
+        min_nbor_dist, sel = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True
+        )
+        local_jdata_cpy["sel"] = sel[0]
+        return local_jdata_cpy, min_nbor_dist
diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py
index 678b797e6c..b33a528721 100644
--- a/deepmd/pt/model/descriptor/dpa2.py
+++ b/deepmd/pt/model/descriptor/dpa2.py
@@ -37,6 +37,9 @@
 from deepmd.pt.utils.utils import (
     to_numpy_array,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -573,30 +576,46 @@ def forward(
         return g1, rot_mat, g2, h2, sw
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
         update_sel = UpdateSel()
-        local_jdata_cpy["repinit"] = update_sel.update_one_sel(
-            global_jdata,
-            local_jdata_cpy["repinit"],
+        min_nbor_dist, repinit_sel = update_sel.update_one_sel(
+            train_data,
+            type_map,
+            local_jdata_cpy["repinit"]["rcut"],
+            local_jdata_cpy["repinit"]["nsel"],
             True,
-            rcut_key="rcut",
-            sel_key="nsel",
         )
-        local_jdata_cpy["repformer"] = update_sel.update_one_sel(
-            global_jdata,
-            local_jdata_cpy["repformer"],
+        local_jdata_cpy["repinit"]["nsel"] = repinit_sel[0]
+        min_nbor_dist, repformer_sel = update_sel.update_one_sel(
+            train_data,
+            type_map,
+            local_jdata_cpy["repformer"]["rcut"],
+            local_jdata_cpy["repformer"]["nsel"],
             True,
-            rcut_key="rcut",
-            sel_key="nsel",
         )
-        return local_jdata_cpy
+        local_jdata_cpy["repformer"]["nsel"] = repformer_sel[0]
+        return local_jdata_cpy, min_nbor_dist
diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py
index bce52fd8ca..e202005f4c 100644
--- a/deepmd/pt/model/descriptor/hybrid.py
+++ b/deepmd/pt/model/descriptor/hybrid.py
@@ -5,6 +5,7 @@
     Dict,
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -20,6 +21,9 @@
 from deepmd.pt.utils.utils import (
     to_torch_tensor,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -242,22 +246,42 @@ def forward(
         return out_descriptor, out_gr, out_g2, out_h2, out_sw
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict:
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        local_jdata_cpy["list"] = [
-            BaseDescriptor.update_sel(global_jdata, sub_jdata)
-            for sub_jdata in local_jdata["list"]
-        ]
-        return local_jdata_cpy
+        new_list = []
+        min_nbor_dist = None
+        for sub_jdata in local_jdata["list"]:
+            new_sub_jdata, min_nbor_dist_ = BaseDescriptor.update_sel(
+                train_data, type_map, sub_jdata
+            )
+            if min_nbor_dist_ is not None:
+                min_nbor_dist = min_nbor_dist_
+            new_list.append(new_sub_jdata)
+        local_jdata_cpy["list"] = new_list
+        return local_jdata_cpy, min_nbor_dist
 
     def serialize(self) -> dict:
         return {
diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
index 0054360593..350fceae2d 100644
--- a/deepmd/pt/model/descriptor/se_a.py
+++ b/deepmd/pt/model/descriptor/se_a.py
@@ -30,6 +30,9 @@
 from deepmd.pt.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.env_mat_stat import (
     StatItem,
 )
@@ -301,18 +304,35 @@ def t_cvt(xx):
         return obj
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False)
+        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
+        )
+        return local_jdata_cpy, min_nbor_dist
 
 
 @DescriptorBlock.register("se_e2_a")
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
index 340fd0c02a..b0a739f5e6 100644
--- a/deepmd/pt/model/descriptor/se_r.py
+++ b/deepmd/pt/model/descriptor/se_r.py
@@ -35,6 +35,9 @@
 from deepmd.pt.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.env_mat_stat import (
     StatItem,
 )
@@ -413,15 +416,32 @@ def t_cvt(xx):
         return obj
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False)
+        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
+        )
+        return local_jdata_cpy, min_nbor_dist
diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py
index db6244000d..2c8f52709f 100644
--- a/deepmd/pt/model/descriptor/se_t.py
+++ b/deepmd/pt/model/descriptor/se_t.py
@@ -30,6 +30,9 @@
 from deepmd.pt.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.env_mat_stat import (
     StatItem,
 )
@@ -324,18 +327,35 @@ def t_cvt(xx):
         return obj
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False)
+        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
+        )
+        return local_jdata_cpy, min_nbor_dist
 
 
 @DescriptorBlock.register("se_e3")
diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py
index fab1ff580f..d3a65db287 100644
--- a/deepmd/pt/model/model/dp_model.py
+++ b/deepmd/pt/model/model/dp_model.py
@@ -1,28 +1,51 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    List,
+    Optional,
+    Tuple,
+)
+
 from deepmd.pt.model.descriptor.base_descriptor import (
     BaseDescriptor,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 
 class DPModelCommon:
     """A base class to implement common methods for all the Models."""
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel(
-            global_jdata, local_jdata["descriptor"]
+        local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel(
+            train_data, type_map, local_jdata["descriptor"]
         )
-        return local_jdata_cpy
+        return local_jdata_cpy, min_nbor_dist
 
     def get_fitting_net(self):
         """Get the fitting network."""
diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py
index f18e1d097f..45b76d1da6 100644
--- a/deepmd/pt/model/model/dp_zbl_model.py
+++ b/deepmd/pt/model/model/dp_zbl_model.py
@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
     Dict,
+    List,
     Optional,
+    Tuple,
 )
 
 import torch
@@ -12,6 +14,9 @@
 from deepmd.pt.model.model.model import (
     BaseModel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 from .dp_model import (
     DPModelCommon,
@@ -105,18 +110,32 @@ def forward_lower(
         return model_predict
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        local_jdata_cpy["dpmodel"] = DPModelCommon.update_sel(
-            global_jdata, local_jdata["dpmodel"]
+        local_jdata_cpy["dpmodel"], min_nbor_dist = DPModelCommon.update_sel(
+            train_data, type_map, local_jdata["dpmodel"]
         )
-        return local_jdata_cpy
+        return local_jdata_cpy, min_nbor_dist
diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py
index e3dcd389bb..148ffaa703 100644
--- a/deepmd/pt/model/model/frozen.py
+++ b/deepmd/pt/model/model/frozen.py
@@ -5,6 +5,7 @@
     Dict,
     List,
     Optional,
+    Tuple,
 )
 
 import torch
@@ -18,6 +19,9 @@
 from deepmd.pt.model.model.model import (
     BaseModel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 
 @BaseModel.register("frozen")
@@ -130,6 +134,11 @@ def get_model_def_script(self) -> str:
         # be a problem
         return self.model.get_model_def_script()
 
+    @torch.jit.export
+    def get_min_nbor_dist(self) -> Optional[float]:
+        """Get the minimum neighbor distance."""
+        return self.model.get_min_nbor_dist()
+
     def serialize(self) -> dict:
         from deepmd.pt.model.model import (
             get_model,
@@ -156,17 +165,31 @@ def get_nsel(self) -> int:
         return self.model.get_nsel()
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
-        return local_jdata
+        return local_jdata, None
 
     @torch.jit.export
     def model_output_type(self) -> str:
diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py
index bf97472e33..d3670737ba 100644
--- a/deepmd/pt/model/model/model.py
+++ b/deepmd/pt/model/model/model.py
@@ -18,6 +18,7 @@ def __init__(self, *args, **kwargs):
         """Construct a basic model for different tasks."""
         torch.nn.Module.__init__(self)
         self.model_def_script = ""
+        self.min_nbor_dist = None
 
     def compute_or_load_stat(
         self,
@@ -46,6 +47,11 @@ def get_model_def_script(self) -> str:
         """Get the model definition script."""
         return self.model_def_script
 
+    @torch.jit.export
+    def get_min_nbor_dist(self) -> Optional[float]:
+        """Get the minimum distance between two atoms."""
+        return self.min_nbor_dist
+
     @torch.jit.export
     def get_ntypes(self):
         """Returns the number of element types."""
diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py
index ff932bc402..369a413f50 100644
--- a/deepmd/pt/model/model/spin_model.py
+++ b/deepmd/pt/model/model/spin_model.py
@@ -291,6 +291,11 @@ def get_model_def_script(self) -> str:
         """Get the model definition script."""
         return self.backbone_model.get_model_def_script()
 
+    @torch.jit.export
+    def get_min_nbor_dist(self) -> Optional[float]:
+        """Get the minimum neighbor distance."""
+        return self.backbone_model.get_min_nbor_dist()
+
     @torch.jit.export
     def get_nnei(self) -> int:
         """Returns the total number of selected neighboring atoms in the cut-off radius."""
diff --git a/deepmd/pt/utils/serialization.py b/deepmd/pt/utils/serialization.py
index c99ddbb3c6..21a2a3fbda 100644
--- a/deepmd/pt/utils/serialization.py
+++ b/deepmd/pt/utils/serialization.py
@@ -50,9 +50,10 @@ def serialize_from_file(model_file: str) -> dict:
         "pt_version": torch.__version__,
         "model": model_dict,
         "model_def_script": model_def_script,
-        # TODO
         "@variables": {},
     }
+    if model.get_min_nbor_dist() is not None:
+        data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist()
     return data
 
 
@@ -72,4 +73,6 @@ def deserialize_to_file(model_file: str, data: dict) -> None:
     # JIT will happy in this way...
     model.model_def_script = json.dumps(data["model_def_script"])
     model = torch.jit.script(model)
+    if "min_nbor_dist" in data.get("@variables", {}):
+        model.min_nbor_dist = data["@variables"]["min_nbor_dist"]
     torch.jit.save(model, model_file)
diff --git a/deepmd/pt/utils/update_sel.py b/deepmd/pt/utils/update_sel.py
index 8c2d0699f2..7f42a9f91c 100644
--- a/deepmd/pt/utils/update_sel.py
+++ b/deepmd/pt/utils/update_sel.py
@@ -15,7 +15,3 @@ class UpdateSel(BaseUpdateSel):
     @property
     def neighbor_stat(self) -> Type[NeighborStat]:
         return NeighborStat
-
-    def hook(self, min_nbor_dist, max_nbor_size):
-        # TODO: save to the model in UpdateSel.hook
-        pass
diff --git a/deepmd/tf/descriptor/descriptor.py b/deepmd/tf/descriptor/descriptor.py
index fabaf78c85..2f813da731 100644
--- a/deepmd/tf/descriptor/descriptor.py
+++ b/deepmd/tf/descriptor/descriptor.py
@@ -26,6 +26,9 @@
 from deepmd.utils.data import (
     DataRequirementItem,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.plugin import (
     make_plugin_registry,
 )
@@ -462,19 +465,33 @@ def explicit_ntypes(self) -> bool:
 
     @classmethod
     @abstractmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         # call subprocess
         cls = cls.get_class_by_type(j_get_type(local_jdata, cls.__name__))
-        return cls.update_sel(global_jdata, local_jdata)
+        return cls.update_sel(train_data, type_map, local_jdata)
 
     @classmethod
     def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor":
diff --git a/deepmd/tf/descriptor/hybrid.py b/deepmd/tf/descriptor/hybrid.py
index 7c69ea5202..fe4fc2ae6a 100644
--- a/deepmd/tf/descriptor/hybrid.py
+++ b/deepmd/tf/descriptor/hybrid.py
@@ -17,6 +17,9 @@
 from deepmd.tf.utils.spin import (
     Spin,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -423,22 +426,42 @@ def explicit_ntypes(self) -> bool:
         return any(ii.explicit_ntypes for ii in self.descrpt_list)
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        local_jdata_cpy["list"] = [
-            Descriptor.update_sel(global_jdata, sub_jdata)
-            for sub_jdata in local_jdata["list"]
-        ]
-        return local_jdata_cpy
+        new_list = []
+        min_nbor_dist = None
+        for sub_jdata in local_jdata["list"]:
+            new_sub_jdata, min_nbor_dist_ = Descriptor.update_sel(
+                train_data, type_map, sub_jdata
+            )
+            if min_nbor_dist_ is not None:
+                min_nbor_dist = min_nbor_dist_
+            new_list.append(new_sub_jdata)
+        local_jdata_cpy["list"] = new_list
+        return local_jdata_cpy, min_nbor_dist
 
     def serialize(self, suffix: str = "") -> dict:
         return {
diff --git a/deepmd/tf/descriptor/loc_frame.py b/deepmd/tf/descriptor/loc_frame.py
index 963e9bf607..e6247290c0 100644
--- a/deepmd/tf/descriptor/loc_frame.py
+++ b/deepmd/tf/descriptor/loc_frame.py
@@ -20,6 +20,9 @@
 from deepmd.tf.utils.sess import (
     run_sess,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 from .descriptor import (
     Descriptor,
@@ -431,14 +434,28 @@ def init_variables(
         self.dstd = get_tensor_by_name_from_graph(graph, f"descrpt_attr{suffix}/t_std")
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
-        return local_jdata
+        return local_jdata, None
diff --git a/deepmd/tf/descriptor/se.py b/deepmd/tf/descriptor/se.py
index a14dc0a5ce..f5f54550f2 100644
--- a/deepmd/tf/descriptor/se.py
+++ b/deepmd/tf/descriptor/se.py
@@ -2,6 +2,7 @@
 import re
 from typing import (
     List,
+    Optional,
     Set,
     Tuple,
 )
@@ -21,6 +22,9 @@
 from deepmd.tf.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 from .descriptor import (
     Descriptor,
@@ -150,19 +154,36 @@ def precision(self) -> tf.DType:
         return self.filter_precision
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         # default behavior is to update sel which is a list
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False)
+        min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False
+        )
+        return local_jdata_cpy, min_nbor_dist
 
     def serialize_network(
         self,
diff --git a/deepmd/tf/descriptor/se_a_mask.py b/deepmd/tf/descriptor/se_a_mask.py
index e78dfba461..b79e806fca 100644
--- a/deepmd/tf/descriptor/se_a_mask.py
+++ b/deepmd/tf/descriptor/se_a_mask.py
@@ -24,6 +24,9 @@
 from deepmd.tf.utils.network import (
     embedding_net_rand_seed_shift,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 from .descriptor import (
     Descriptor,
@@ -420,14 +423,28 @@ def prod_force_virial(
         return force, virial, atom_virial
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
-        return local_jdata
+        return local_jdata, None
diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py
index 43c38b0955..6d3cfeaa6e 100644
--- a/deepmd/tf/descriptor/se_atten.py
+++ b/deepmd/tf/descriptor/se_atten.py
@@ -78,6 +78,9 @@
 from deepmd.tf.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 from deepmd.utils.version import (
     check_version_compatibility,
 )
@@ -1470,18 +1473,36 @@ def explicit_ntypes(self) -> bool:
         return True
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True)
+        min_nbor_dist, sel = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True
+        )
+        local_jdata_cpy["sel"] = sel[0]
+        return local_jdata_cpy, min_nbor_dist
 
     def serialize_attention_layers(
         self,
diff --git a/deepmd/tf/entrypoints/compress.py b/deepmd/tf/entrypoints/compress.py
index 1e50d88b63..6f3ff13e3b 100644
--- a/deepmd/tf/entrypoints/compress.py
+++ b/deepmd/tf/entrypoints/compress.py
@@ -32,6 +32,9 @@
 from deepmd.tf.utils.update_sel import (
     UpdateSel,
 )
+from deepmd.utils.data_system import (
+    get_data,
+)
 
 from .freeze import (
     freeze,
@@ -115,9 +118,17 @@ def compress(
             log.info("stage 0: compute the min_nbor_dist")
             jdata = j_loader(training_script)
             jdata = update_deepmd_input(jdata)
+
+            type_map = jdata["model"].get("type_map", None)
+            train_data = get_data(
+                jdata["training"]["training_data"],
+                0,  # not used
+                type_map,
+                None,
+            )
             update_sel = UpdateSel()
             t_min_nbor_dist = update_sel.get_min_nbor_dist(
-                jdata, update_sel.get_rcut(jdata)
+                train_data,
             )
 
     _check_compress_type(graph)
diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py
index 3c4decbe8c..d394773cf2 100755
--- a/deepmd/tf/entrypoints/train.py
+++ b/deepmd/tf/entrypoints/train.py
@@ -17,6 +17,7 @@
     j_loader,
 )
 from deepmd.tf.env import (
+    GLOBAL_ENER_FLOAT_PRECISION,
     reset_default_tf_session_config,
     tf,
 )
@@ -255,5 +256,21 @@ def update_sel(jdata):
         "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)"
     )
     jdata_cpy = jdata.copy()
-    jdata_cpy["model"] = Model.update_sel(jdata, jdata["model"])
+    type_map = jdata["model"].get("type_map")
+    train_data = get_data(
+        jdata["training"]["training_data"],
+        0,  # not used
+        type_map,
+        None,  # not used
+    )
+    jdata_cpy["model"], min_nbor_dist = Model.update_sel(
+        train_data, type_map, jdata["model"]
+    )
+
+    if min_nbor_dist is not None:
+        tf.constant(
+            min_nbor_dist,
+            name="train_attr/min_nbor_dist",
+            dtype=GLOBAL_ENER_FLOAT_PRECISION,
+        )
     return jdata_cpy
diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py
index fa28b2cc58..3e296c00f2 100644
--- a/deepmd/tf/model/frozen.py
+++ b/deepmd/tf/model/frozen.py
@@ -8,6 +8,7 @@
 from typing import (
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -38,6 +39,9 @@
 from deepmd.utils.data import (
     DataRequirementItem,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 from .model import (
     Model,
@@ -237,18 +241,32 @@ def get_type_map(self) -> list:
         return self.model.get_type_map()
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         # we don't know how to compress it, so no neighbor statistics here
-        return local_jdata
+        return local_jdata, None
 
     def serialize(self, suffix: str = "") -> dict:
         # try to recover the original model
diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py
index 26bc382569..1bd1644e54 100644
--- a/deepmd/tf/model/linear.py
+++ b/deepmd/tf/model/linear.py
@@ -10,6 +10,7 @@
 from typing import (
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -27,6 +28,9 @@
 from deepmd.utils.data import (
     DataRequirementItem,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 from .model import (
     Model,
@@ -133,22 +137,42 @@ def get_type_map(self) -> list:
         return self.models[0].get_type_map()
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        local_jdata_cpy["models"] = [
-            Model.update_sel(global_jdata, sub_jdata)
-            for sub_jdata in local_jdata["models"]
-        ]
-        return local_jdata_cpy
+        new_list = []
+        min_nbor_dist = None
+        for sub_jdata in local_jdata["models"]:
+            new_sub_jdata, min_nbor_dist_ = Model.update_sel(
+                train_data, type_map, sub_jdata
+            )
+            if min_nbor_dist_ is not None:
+                min_nbor_dist = min_nbor_dist_
+            new_list.append(new_sub_jdata)
+        local_jdata_cpy["models"] = new_list
+        return local_jdata_cpy, min_nbor_dist
 
     @property
     def input_requirement(self) -> List[DataRequirementItem]:
diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py
index 194750a9d7..a1baf85dbc 100644
--- a/deepmd/tf/model/model.py
+++ b/deepmd/tf/model/model.py
@@ -11,6 +11,7 @@
     Dict,
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -511,7 +512,12 @@ def get_feed_dict(
 
     @classmethod
     @abstractmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict:
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Notes
@@ -520,8 +526,10 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict:
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
 
@@ -529,9 +537,11 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict:
         -------
         dict
             The updated local data
+        float
+            The minimum distance between two atoms
         """
         cls = cls.get_class_by_type(local_jdata.get("type", "standard"))
-        return cls.update_sel(global_jdata, local_jdata)
+        return cls.update_sel(train_data, type_map, local_jdata)
 
     @classmethod
     def deserialize(cls, data: dict, suffix: str = "") -> "Model":
@@ -744,21 +754,35 @@ def get_ntypes(self) -> int:
         return self.ntypes
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        local_jdata_cpy["descriptor"] = Descriptor.update_sel(
-            global_jdata, local_jdata["descriptor"]
+        local_jdata_cpy["descriptor"], min_nbor_dist = Descriptor.update_sel(
+            train_data, type_map, local_jdata["descriptor"]
         )
-        return local_jdata_cpy
+        return local_jdata_cpy, min_nbor_dist
 
     @classmethod
     def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor":
diff --git a/deepmd/tf/model/pairtab.py b/deepmd/tf/model/pairtab.py
index 4f71dcd76e..8c79dedea0 100644
--- a/deepmd/tf/model/pairtab.py
+++ b/deepmd/tf/model/pairtab.py
@@ -5,6 +5,7 @@
 from typing import (
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -35,6 +36,9 @@
 from deepmd.utils.data import (
     DataRequirementItem,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 
 @Model.register("pairtab")
@@ -268,7 +272,12 @@ def enable_compression(self, suffix: str = "") -> None:
         # nothing needs to do
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict:
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Notes
@@ -277,8 +286,10 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict:
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
 
@@ -286,9 +297,15 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict:
         -------
         dict
             The updated local data
+        float
+            The minimum distance between two atoms
         """
         local_jdata_cpy = local_jdata.copy()
-        return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True)
+        min_nbor_dist, sel = UpdateSel().update_one_sel(
+            train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True
+        )
+        local_jdata_cpy["sel"] = sel[0]
+        return local_jdata_cpy, min_nbor_dist
 
     @property
     def input_requirement(self) -> List[DataRequirementItem]:
diff --git a/deepmd/tf/model/pairwise_dprc.py b/deepmd/tf/model/pairwise_dprc.py
index 3d61dfd339..44e3943e12 100644
--- a/deepmd/tf/model/pairwise_dprc.py
+++ b/deepmd/tf/model/pairwise_dprc.py
@@ -3,6 +3,7 @@
     Dict,
     List,
     Optional,
+    Tuple,
     Union,
 )
 
@@ -36,6 +37,9 @@
 from deepmd.utils.data import (
     DataRequirementItem,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 
 @Model.register("pairwise_dprc")
@@ -410,20 +414,33 @@ def get_feed_dict(
         return feed_dict
 
     @classmethod
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ) -> Tuple[dict, Optional[float]]:
         """Update the selection and perform neighbor statistics.
 
         Parameters
         ----------
-        global_jdata : dict
-            The global data, containing the training section
+        train_data : DeepmdDataSystem
+            data used to do neighbor statictics
+        type_map : list[str], optional
+            The name of each type of atoms
         local_jdata : dict
             The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
         """
         # do not update sel; only find min distance
-        # rcut is not important here
-        UpdateSel().get_min_nbor_dist(global_jdata, 6.0)
-        return local_jdata
+        min_nbor_dist = UpdateSel().get_min_nbor_dist(train_data)
+        return local_jdata, min_nbor_dist
 
     @property
     def input_requirement(self) -> List[DataRequirementItem]:
diff --git a/deepmd/tf/utils/update_sel.py b/deepmd/tf/utils/update_sel.py
index db0420dde8..726aec4d41 100644
--- a/deepmd/tf/utils/update_sel.py
+++ b/deepmd/tf/utils/update_sel.py
@@ -3,12 +3,6 @@
     Type,
 )
 
-from deepmd.env import (
-    GLOBAL_ENER_FLOAT_PRECISION,
-)
-from deepmd.tf.env import (
-    tf,
-)
 from deepmd.tf.utils.neighbor_stat import (
     NeighborStat,
 )
@@ -21,12 +15,3 @@ class UpdateSel(BaseUpdateSel):
     @property
     def neighbor_stat(self) -> Type[NeighborStat]:
         return NeighborStat
-
-    def hook(self, min_nbor_dist, max_nbor_size):
-        # moved from traier.py as duplicated
-        tf.constant(
-            min_nbor_dist,
-            name="train_attr/min_nbor_dist",
-            dtype=GLOBAL_ENER_FLOAT_PRECISION,
-        )
-        tf.constant(max_nbor_size, name="train_attr/max_nbor_size", dtype=tf.int32)
diff --git a/deepmd/utils/update_sel.py b/deepmd/utils/update_sel.py
index d1be8e8138..6feed525e5 100644
--- a/deepmd/utils/update_sel.py
+++ b/deepmd/utils/update_sel.py
@@ -1,14 +1,19 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
 from abc import (
+    ABC,
     abstractmethod,
 )
 from typing import (
+    List,
+    Optional,
+    Tuple,
     Type,
+    Union,
 )
 
 from deepmd.utils.data_system import (
-    get_data,
+    DeepmdDataSystem,
 )
 from deepmd.utils.neighbor_stat import (
     NeighborStat,
@@ -17,32 +22,29 @@
 log = logging.getLogger(__name__)
 
 
-class BaseUpdateSel:
+class BaseUpdateSel(ABC):
     """Update the sel field in the descriptor."""
 
     def update_one_sel(
         self,
-        jdata,
-        descriptor,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        rcut: float,
+        sel: Union[int, List[int], str],
         mixed_type: bool = False,
-        rcut_key="rcut",
-        sel_key="sel",
-    ):
-        rcut = descriptor[rcut_key]
-        tmp_sel = self.get_sel(
-            jdata,
+    ) -> Tuple[float, List[int]]:
+        min_nbor_dist, tmp_sel = self.get_nbor_stat(
+            train_data,
+            type_map,
             rcut,
             mixed_type=mixed_type,
         )
-        sel = descriptor[sel_key]
         if isinstance(sel, int):
             # convert to list and finnally convert back to int
             sel = [sel]
-        if self.parse_auto_sel(descriptor[sel_key]):
-            ratio = self.parse_auto_sel_ratio(descriptor[sel_key])
-            descriptor[sel_key] = sel = [
-                int(self.wrap_up_4(ii * ratio)) for ii in tmp_sel
-            ]
+        if self.parse_auto_sel(sel):
+            ratio = self.parse_auto_sel_ratio(sel)
+            sel = [int(self.wrap_up_4(ii * ratio)) for ii in tmp_sel]
         else:
             # sel is set by user
             for ii, (tt, dd) in enumerate(zip(tmp_sel, sel)):
@@ -54,9 +56,7 @@ def update_one_sel(
                         "not less than %d, but you set it to %d. The accuracy"
                         " of your model may get worse." % (ii, tt, dd)
                     )
-        if mixed_type:
-            descriptor[sel_key] = sum(sel)
-        return descriptor
+        return min_nbor_dist, sel
 
     def parse_auto_sel(self, sel):
         if not isinstance(sel, str):
@@ -83,65 +83,36 @@ def parse_auto_sel_ratio(self, sel):
     def wrap_up_4(self, xx):
         return 4 * ((int(xx) + 3) // 4)
 
-    def get_sel(self, jdata, rcut, mixed_type: bool = False):
-        _, max_nbor_size = self.get_nbor_stat(jdata, rcut, mixed_type=mixed_type)
-        return max_nbor_size
-
-    def get_rcut(self, jdata):
-        if jdata["model"].get("type") == "pairwise_dprc":
-            return max(
-                jdata["model"]["qm_model"]["descriptor"]["rcut"],
-                jdata["model"]["qmmm_model"]["descriptor"]["rcut"],
-            )
-        descrpt_data = jdata["model"]["descriptor"]
-        rcut_list = []
-        if descrpt_data["type"] == "hybrid":
-            for ii in descrpt_data["list"]:
-                rcut_list.append(ii["rcut"])
-        else:
-            rcut_list.append(descrpt_data["rcut"])
-        return max(rcut_list)
-
-    def get_type_map(self, jdata):
-        return jdata["model"].get("type_map", None)
-
-    def get_nbor_stat(self, jdata, rcut, mixed_type: bool = False):
-        # it seems that DeepmdDataSystem does not need rcut
-        # it's not clear why there is an argument...
-        # max_rcut = get_rcut(jdata)
-        max_rcut = rcut
-        type_map = self.get_type_map(jdata)
-
+    def get_nbor_stat(
+        self,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        rcut: float,
+        mixed_type: bool = False,
+    ) -> Tuple[float, Union[int, List[int]]]:
+        """Get the neighbor statistics of the data.
+
+        Parameters
+        ----------
+        train_data : DeepmdDataSystem
+            The training data.
+        type_map : Optional[List[str]]
+            The type map.
+        rcut : float
+            The cutoff radius.
+        mixed_type : bool, optional
+            Whether to mix the types.
+
+        Returns
+        -------
+        min_nbor_dist : float
+            The minimum neighbor distance.
+        max_nbor_size : List[int]
+            The maximum neighbor size.
+        """
         if type_map and len(type_map) == 0:
             type_map = None
-        multi_task_mode = "data_dict" in jdata["training"]
-        if not multi_task_mode:
-            train_data = get_data(
-                jdata["training"]["training_data"], max_rcut, type_map, None
-            )
-            train_data.get_batch()
-        else:
-            assert (
-                type_map is not None
-            ), "Data stat in multi-task mode must have available type_map! "
-            train_data = None
-            for systems in jdata["training"]["data_dict"]:
-                tmp_data = get_data(
-                    jdata["training"]["data_dict"][systems]["training_data"],
-                    max_rcut,
-                    type_map,
-                    None,
-                )
-                tmp_data.get_batch()
-                assert tmp_data.get_type_map(), f"In multi-task mode, 'type_map.raw' must be defined in data systems {systems}! "
-                if train_data is None:
-                    train_data = tmp_data
-                else:
-                    train_data.system_dirs += tmp_data.system_dirs
-                    train_data.data_systems += tmp_data.data_systems
-                    train_data.natoms += tmp_data.natoms
-                    train_data.natoms_vec += tmp_data.natoms_vec
-                    train_data.default_mesh += tmp_data.default_mesh
+        train_data.get_batch()
         data_ntypes = train_data.get_ntypes()
         if type_map is not None:
             map_ntypes = len(type_map)
@@ -152,7 +123,6 @@ def get_nbor_stat(self, jdata, rcut, mixed_type: bool = False):
         neistat = self.neighbor_stat(ntypes, rcut, mixed_type=mixed_type)
 
         min_nbor_dist, max_nbor_size = neistat.get_stat(train_data)
-        self.hook(min_nbor_dist, max_nbor_size)
 
         return min_nbor_dist, max_nbor_size
 
@@ -161,10 +131,14 @@ def get_nbor_stat(self, jdata, rcut, mixed_type: bool = False):
     def neighbor_stat(self) -> Type[NeighborStat]:
         pass
 
-    @abstractmethod
-    def hook(self, min_nbor_dist, max_nbor_size):
-        pass
-
-    def get_min_nbor_dist(self, jdata, rcut):
-        min_nbor_dist, _ = self.get_nbor_stat(jdata, rcut)
+    def get_min_nbor_dist(
+        self,
+        train_data: DeepmdDataSystem,
+    ):
+        min_nbor_dist, _ = self.get_nbor_stat(
+            train_data,
+            None,  # type_map doesn't affect min_nbor_dist
+            1e-6,  # we don't need the max_nbor_size
+            mixed_type=True,  # mixed_types doesn't affect min_nbor_dist
+        )
         return min_nbor_dist
diff --git a/doc/development/create-a-model-pt.md b/doc/development/create-a-model-pt.md
index 35d81b364a..c6d372b195 100644
--- a/doc/development/create-a-model-pt.md
+++ b/doc/development/create-a-model-pt.md
@@ -20,6 +20,7 @@ The framework-independent backend is implemented in pure NumPy, serving as a ref
 When creating a new descriptor, it is essential to inherit from both the {py:class}`deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor` class and the {py:class}`torch.nn.Module` class. Abstract methods, including {py:class}`deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor.forward`, must be implemented, while others remain optional. It is crucial to adhere to the original method arguments without any modifications. Once the implementation is complete, the next step involves registering the component with a designated key:
 
 ```py
+from deepmd.utils.data_system import DeepmdDataSystem
 from deepmd.pt.model.descriptor.base_descriptor import (
     BaseDescriptor,
 )
@@ -63,7 +64,12 @@ class SomeDescript(BaseDescriptor, torch.nn.Module):
     def deserialize(cls, data: dict) -> "SomeDescript":
         pass
 
-    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: Optional[List[str]],
+        local_jdata: dict,
+    ):
         pass
 ```
 
diff --git a/source/tests/common/dpmodel/test_update_sel.py b/source/tests/common/dpmodel/test_update_sel.py
new file mode 100644
index 0000000000..a2656da7de
--- /dev/null
+++ b/source/tests/common/dpmodel/test_update_sel.py
@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import random
+import unittest
+from unittest.mock import (
+    patch,
+)
+
+from deepmd.dpmodel.model.base_model import (
+    BaseModel,
+)
+from deepmd.dpmodel.utils.update_sel import (
+    UpdateSel,
+)
+
+
+def update_sel(jdata):
+    type_map = jdata["model"].get("type_map")
+    train_data = None
+    jdata["model"], _ = BaseModel.update_sel(train_data, type_map, jdata["model"])
+    return jdata
+
+
+class TestTrain(unittest.TestCase):
+    def setUp(self) -> None:
+        self.update_sel = UpdateSel()
+        self.mock_min_nbor_dist = random.random()
+        return super().setUp()
+
+    @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_one_sel(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [10, 20]
+
+        min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto")
+        # self.assertEqual(descriptor['sel'], [11,22])
+        self.assertEqual(sel, [12, 24])
+        self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist)
+        min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto:1.5")
+        # self.assertEqual(descriptor['sel'], [15,30])
+        self.assertEqual(sel, [16, 32])
+        self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist)
+
+    @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_hybrid(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [10, 20]
+
+        jdata = {
+            "model": {
+                "descriptor": {
+                    "type": "hybrid",
+                    "list": [
+                        {"type": "se_e2_a", "rcut": 6, "sel": "auto"},
+                        {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"},
+                    ],
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {
+                "descriptor": {
+                    "type": "hybrid",
+                    "list": [
+                        {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]},
+                        {"type": "se_e2_a", "rcut": 6, "sel": [16, 32]},
+                    ],
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [10, 20]
+
+        jdata = {
+            "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}},
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}},
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_atten_auto(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+
+        jdata = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": "auto",
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 28,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_atten_int(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+
+        jdata = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 30,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 30,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_atten_list(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+
+        jdata = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 30,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 30,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    def test_wrap_up_4(self):
+        self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(14), 4 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(15), 4 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(16), 4 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(17), 5 * 4)
diff --git a/source/tests/pt/test_update_sel.py b/source/tests/pt/test_update_sel.py
new file mode 100644
index 0000000000..e4eff0ebf3
--- /dev/null
+++ b/source/tests/pt/test_update_sel.py
@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import random
+import unittest
+from unittest.mock import (
+    patch,
+)
+
+from deepmd.pt.model.model.model import (
+    BaseModel,
+)
+from deepmd.pt.utils.update_sel import (
+    UpdateSel,
+)
+
+
+def update_sel(jdata):
+    type_map = jdata["model"].get("type_map")
+    train_data = None
+    jdata["model"], _ = BaseModel.update_sel(train_data, type_map, jdata["model"])
+    return jdata
+
+
+class TestTrain(unittest.TestCase):
+    def setUp(self) -> None:
+        self.update_sel = UpdateSel()
+        self.mock_min_nbor_dist = random.random()
+        return super().setUp()
+
+    @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_one_sel(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [10, 20]
+
+        min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto")
+        # self.assertEqual(descriptor['sel'], [11,22])
+        self.assertEqual(sel, [12, 24])
+        self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist)
+        min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto:1.5")
+        # self.assertEqual(descriptor['sel'], [15,30])
+        self.assertEqual(sel, [16, 32])
+        self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist)
+
+    @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_hybrid(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [10, 20]
+
+        jdata = {
+            "model": {
+                "descriptor": {
+                    "type": "hybrid",
+                    "list": [
+                        {"type": "se_e2_a", "rcut": 6, "sel": "auto"},
+                        {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"},
+                    ],
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {
+                "descriptor": {
+                    "type": "hybrid",
+                    "list": [
+                        {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]},
+                        {"type": "se_e2_a", "rcut": 6, "sel": [16, 32]},
+                    ],
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [10, 20]
+
+        jdata = {
+            "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}},
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}},
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_atten_auto(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+
+        jdata = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": "auto",
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 28,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_atten_int(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+
+        jdata = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 30,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 30,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_atten_list(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+
+        jdata = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 30,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {
+                "descriptor": {
+                    "type": "se_atten",
+                    "sel": 30,
+                    "rcut": 6,
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    def test_skip_frozen(self):
+        jdata = {
+            "model": {
+                "type": "frozen",
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = jdata.copy()
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
+    def test_wrap_up_4(self):
+        self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(14), 4 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(15), 4 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(16), 4 * 4)
+        self.assertEqual(self.update_sel.wrap_up_4(17), 5 * 4)
diff --git a/source/tests/tf/test_train.py b/source/tests/tf/test_train.py
index 3e22dc57bc..f5a57a948d 100644
--- a/source/tests/tf/test_train.py
+++ b/source/tests/tf/test_train.py
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import random
 import unittest
 from unittest.mock import (
     patch,
@@ -15,6 +16,7 @@
 class TestTrain(unittest.TestCase):
     def setUp(self) -> None:
         self.update_sel = UpdateSel()
+        self.mock_min_nbor_dist = random.random()
         return super().setUp()
 
     def test_train_parse_auto_sel(self):
@@ -34,22 +36,25 @@ def test_train_parse_auto_sel_ratio(self):
         with self.assertRaises(RuntimeError):
             self.update_sel.parse_auto_sel_ratio([1, 2, 3])
 
-    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel")
-    def test_update_one_sel(self, sel_mock):
-        sel_mock.return_value = [10, 20]
-        jdata = {}
-        descriptor = {"type": "se_e2_a", "rcut": 6, "sel": "auto"}
-        descriptor = self.update_sel.update_one_sel(jdata, descriptor)
+    @patch("deepmd.tf.entrypoints.train.get_data")
+    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_one_sel(self, sel_mock, get_data_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [10, 20]
+        get_data_mock.return_value = None
+        min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto")
         # self.assertEqual(descriptor['sel'], [11,22])
-        self.assertEqual(descriptor["sel"], [12, 24])
-        descriptor = {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"}
-        descriptor = self.update_sel.update_one_sel(jdata, descriptor)
+        self.assertEqual(sel, [12, 24])
+        self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist)
+        min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto:1.5")
         # self.assertEqual(descriptor['sel'], [15,30])
-        self.assertEqual(descriptor["sel"], [16, 32])
+        self.assertEqual(sel, [16, 32])
+        self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist)
 
-    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel")
-    def test_update_sel_hybrid(self, sel_mock):
-        sel_mock.return_value = [10, 20]
+    @patch("deepmd.tf.entrypoints.train.get_data")
+    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_hybrid(self, sel_mock, get_data_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [10, 20]
+        get_data_mock.return_value = None
         jdata = {
             "model": {
                 "descriptor": {
@@ -59,7 +64,8 @@ def test_update_sel_hybrid(self, sel_mock):
                         {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"},
                     ],
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         expected_out = {
             "model": {
@@ -70,24 +76,33 @@ def test_update_sel_hybrid(self, sel_mock):
                         {"type": "se_e2_a", "rcut": 6, "sel": [16, 32]},
                     ],
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel")
-    def test_update_sel(self, sel_mock):
-        sel_mock.return_value = [10, 20]
-        jdata = {"model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}}}
+    @patch("deepmd.tf.entrypoints.train.get_data")
+    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel(self, sel_mock, get_data_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [10, 20]
+        get_data_mock.return_value = None
+        jdata = {
+            "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}},
+            "training": {"training_data": {}},
+        }
         expected_out = {
-            "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}}
+            "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}},
+            "training": {"training_data": {}},
         }
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel")
-    def test_update_sel_atten_auto(self, sel_mock):
-        sel_mock.return_value = [25]
+    @patch("deepmd.tf.entrypoints.train.get_data")
+    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_atten_auto(self, sel_mock, get_data_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+        get_data_mock.return_value = None
         jdata = {
             "model": {
                 "descriptor": {
@@ -95,7 +110,8 @@ def test_update_sel_atten_auto(self, sel_mock):
                     "sel": "auto",
                     "rcut": 6,
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         expected_out = {
             "model": {
@@ -104,14 +120,17 @@ def test_update_sel_atten_auto(self, sel_mock):
                     "sel": 28,
                     "rcut": 6,
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel")
-    def test_update_sel_atten_int(self, sel_mock):
-        sel_mock.return_value = [25]
+    @patch("deepmd.tf.entrypoints.train.get_data")
+    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_atten_int(self, sel_mock, get_data_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+        get_data_mock.return_value = None
         jdata = {
             "model": {
                 "descriptor": {
@@ -119,7 +138,8 @@ def test_update_sel_atten_int(self, sel_mock):
                     "sel": 30,
                     "rcut": 6,
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         expected_out = {
             "model": {
@@ -128,14 +148,17 @@ def test_update_sel_atten_int(self, sel_mock):
                     "sel": 30,
                     "rcut": 6,
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel")
-    def test_update_sel_atten_list(self, sel_mock):
-        sel_mock.return_value = [25]
+    @patch("deepmd.tf.entrypoints.train.get_data")
+    @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_atten_list(self, sel_mock, get_data_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+        get_data_mock.return_value = None
         jdata = {
             "model": {
                 "descriptor": {
@@ -143,7 +166,8 @@ def test_update_sel_atten_list(self, sel_mock):
                     "sel": 30,
                     "rcut": 6,
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         expected_out = {
             "model": {
@@ -152,19 +176,23 @@ def test_update_sel_atten_list(self, sel_mock):
                     "sel": 30,
                     "rcut": 6,
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    def test_skip_loc_frame(self):
+    @patch("deepmd.tf.entrypoints.train.get_data")
+    def test_skip_loc_frame(self, get_data_mock):
+        get_data_mock.return_value = None
         jdata = {
             "model": {
                 "descriptor": {
                     "type": "loc_frame",
                     "rcut": 6,
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         expected_out = {
             "model": {
@@ -172,22 +200,28 @@ def test_skip_loc_frame(self):
                     "type": "loc_frame",
                     "rcut": 6,
                 }
-            }
+            },
+            "training": {"training_data": {}},
         }
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    def test_skip_frozen(self):
+    @patch("deepmd.tf.entrypoints.train.get_data")
+    def test_skip_frozen(self, get_data_mock):
+        get_data_mock.return_value = None
         jdata = {
             "model": {
                 "type": "frozen",
-            }
+            },
+            "training": {"training_data": {}},
         }
         expected_out = jdata.copy()
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
-    def test_skip_linear_frozen(self):
+    @patch("deepmd.tf.entrypoints.train.get_data")
+    def test_skip_linear_frozen(self, get_data_mock):
+        get_data_mock.return_value = None
         jdata = {
             "model": {
                 "type": "linear_ener",
@@ -197,15 +231,18 @@ def test_skip_linear_frozen(self):
                     {"type": "frozen"},
                     {"type": "frozen"},
                 ],
-            }
+            },
+            "training": {"training_data": {}},
         }
         expected_out = jdata.copy()
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
+    @patch("deepmd.tf.entrypoints.train.get_data")
     @patch("deepmd.tf.utils.update_sel.UpdateSel.get_min_nbor_dist")
-    def test_pairwise_dprc(self, sel_mock):
-        sel_mock.return_value = 0.5
+    def test_pairwise_dprc(self, sel_mock, get_data_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist
+        get_data_mock.return_value = None
         jdata = {
             "model": {
                 "type": "pairwise_dprc",
@@ -215,7 +252,8 @@ def test_pairwise_dprc(self, sel_mock):
                     {"type": "frozen"},
                     {"type": "frozen"},
                 ],
-            }
+            },
+            "training": {"training_data": {}},
         }
         expected_out = jdata.copy()
         jdata = update_sel(jdata)