diff --git a/README.md b/README.md
index 40dad38b23..d0bae50b6e 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,11 @@
 # DeePMD-kit(PaddlePaddle backend)
 
 > [!IMPORTANT]
-> 本项目为 DeePMD-kit 的 PaddlePaddle 版本，主要修改了部分代码，使其可以运行在 PaddlePaddle 上。运行功能包括 water_se_e2_a 案例的单卡 GPU 训练、单卡 GPU 评估、导出静态图模型、接入 LAMMPS(GPU) 推理 4 部分的功能。
+> 本项目为 DeePMD-kit 的 PaddlePaddle 版本，修改了部分代码，使其可以以 PaddlePaddle(GPU) 为后端进行训练、评估、模型导出、LAMMPS 推理等任务。案例支持情况如下所示。
+> | example | Train | Test | Export | LAMMPS |
+> | :-----: | :--:  | :--: | :----: | :---: |
+> | water/se_e2_a  | ✅    | ✅    | ✅    | ✅    |
+> | spin/se_e2_a  | ✅    | ✅    | ✅     | TODO     |
 
 ## 1. 环境安装
 
@@ -20,9 +24,8 @@
 
 3. 安装 deepmd-kit
 
-
     ``` sh
-    git clone https://github.com/HydrogenSulfate/deepmd-kit.git -b add_ddle_backend_polish_ver
+    git clone https://github.com/deepmodeling/deepmd-kit.git -b paddle2
     cd deepmd-kit
     # 以 editable 的方式安装，方便调试
     pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple
diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py
index 58f1139ffb..d1044d51d1 100644
--- a/deepmd/descriptor/se_a.py
+++ b/deepmd/descriptor/se_a.py
@@ -203,9 +203,17 @@ def __init__(
         self.useBN = False
         self.dstd = None
         self.davg = None
-        self.avg_zero = paddle.zeros([self.ntypes, self.ndescrpt], dtype="float32")
-        self.std_ones = paddle.ones([self.ntypes, self.ndescrpt], dtype="float32")
-
+        # self.compress = False
+        # self.embedding_net_variables = None
+        # self.mixed_prec = None
+        # self.place_holders = {}
+        # self.nei_type = np.repeat(np.arange(self.ntypes), self.sel_a)
+        self.avg_zero = paddle.zeros(
+            [self.ntypes, self.ndescrpt], dtype=GLOBAL_PD_FLOAT_PRECISION
+        )
+        self.std_ones = paddle.ones(
+            [self.ntypes, self.ndescrpt], dtype=GLOBAL_PD_FLOAT_PRECISION
+        )
         nets = []
         for type_input in range(self.ntypes):
             layer = []
@@ -242,11 +250,19 @@ def __init__(
             }
 
         self.t_rcut = paddle.to_tensor(
-            np.max([self.rcut_r, self.rcut_a]), dtype="float32"
+            np.max([self.rcut_r, self.rcut_a]), dtype=GLOBAL_PD_FLOAT_PRECISION
+        )
+        self.register_buffer("buffer_sel", paddle.to_tensor(self.sel_a, dtype="int32"))
+        self.register_buffer(
+            "buffer_ndescrpt", paddle.to_tensor(self.ndescrpt, dtype="int32")
+        )
+        self.register_buffer(
+            "buffer_original_sel",
+            paddle.to_tensor(
+                self.original_sel if self.original_sel is not None else self.sel_a,
+                dtype="int32",
+            ),
         )
-        self.t_ntypes = paddle.to_tensor(self.ntypes, dtype="int32")
-        self.t_ndescrpt = paddle.to_tensor(self.ndescrpt, dtype="int32")
-        self.t_sel = paddle.to_tensor(self.sel_a, dtype="int32")
 
         t_avg = paddle.to_tensor(
             np.zeros([self.ntypes, self.ndescrpt]), dtype="float64"
@@ -539,6 +555,7 @@ def forward(
         coord = paddle.reshape(coord_, [-1, natoms[1] * 3])
         box = paddle.reshape(box_, [-1, 9])
         atype = paddle.reshape(atype_, [-1, natoms[1]])
+
         (
             self.descrpt,
             self.descrpt_deriv,
@@ -669,7 +686,7 @@ def _pass_filter(
                     [0, start_index, 0],
                     [
                         inputs.shape[0],
-                        start_index + natoms[2 + type_i],
+                        start_index + natoms[2 + type_i].item(),
                         inputs.shape[2],
                     ],
                 )
@@ -697,7 +714,7 @@ def _pass_filter(
                 )
                 output.append(layer)
                 output_qmat.append(qmat)
-                start_index += natoms[2 + type_i]
+                start_index += natoms[2 + type_i].item()
         else:
             raise NotImplementedError()
             # This branch will not be excecuted at current
@@ -747,13 +764,11 @@ def _compute_dstats_sys_smth(
         self, data_coord, data_box, data_atype, natoms_vec, mesh
     ):
         input_dict = {}
-        input_dict["coord"] = paddle.to_tensor(data_coord, dtype="float32")
-        input_dict["box"] = paddle.to_tensor(data_box, dtype="float32")
-        input_dict["type"] = paddle.to_tensor(data_atype, dtype="int32")
-        input_dict["natoms_vec"] = paddle.to_tensor(
-            natoms_vec, dtype="int32", place="cpu"
-        )
-        input_dict["default_mesh"] = paddle.to_tensor(mesh, dtype="int32")
+        input_dict["coord"] = paddle.to_tensor(data_coord, GLOBAL_PD_FLOAT_PRECISION)
+        input_dict["box"] = paddle.to_tensor(data_box, GLOBAL_PD_FLOAT_PRECISION)
+        input_dict["type"] = paddle.to_tensor(data_atype, "int32")
+        input_dict["natoms_vec"] = paddle.to_tensor(natoms_vec, "int32", place="cpu")
+        input_dict["default_mesh"] = paddle.to_tensor(mesh, "int32")
 
         self.stat_descrpt, descrpt_deriv, rij, nlist = op_module.prod_env_mat_a(
             input_dict["coord"],
@@ -949,10 +964,8 @@ def _filter_lower(
             # natom x 4 x outputs_size
 
             return paddle.matmul(
-                paddle.reshape(
-                    inputs_i, [natom, shape_i[1] // 4, 4]
-                ),  # [natom, nei_type_i, 4]
-                xyz_scatter_out,  # [natom, nei_type_i, 100]
+                paddle.reshape(inputs_i, [natom, shape_i[1] // 4, 4]),
+                xyz_scatter_out,
                 transpose_x=True,
             )
 
diff --git a/deepmd/entrypoints/freeze.py b/deepmd/entrypoints/freeze.py
index 57efbc09ce..dba1f7bd15 100755
--- a/deepmd/entrypoints/freeze.py
+++ b/deepmd/entrypoints/freeze.py
@@ -350,7 +350,7 @@ def freeze_graph(
         input_spec=[
             InputSpec(shape=[None], dtype="float64"),  # coord_
             InputSpec(shape=[None], dtype="int32"),  # atype_
-            InputSpec(shape=[4], dtype="int32"),  # natoms
+            InputSpec(shape=[2 + dp.model.descrpt.ntypes], dtype="int32"),  # natoms
             InputSpec(shape=[None], dtype="float64"),  # box
             InputSpec(shape=[6], dtype="int32"),  # mesh
             {
@@ -362,9 +362,9 @@ def freeze_graph(
     )
     for name, param in st_model.named_buffers():
         print(
-            f"[{name}, {param.shape}] generated name in static_model is: {param.name}"
+            f"[{name}, {param.dtype}, {param.shape}] generated name in static_model is: {param.name}"
         )
-    #  skip pruning for program so as to keep buffers into files
+    # skip pruning for program so as to keep buffers into files
     skip_prune_program = True
     print(f"==>> Set skip_prune_program = {skip_prune_program}")
     paddle.jit.save(st_model, output, skip_prune_program=skip_prune_program)
diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py
index 27bf6a2105..4fc87b167b 100644
--- a/deepmd/fit/ener.py
+++ b/deepmd/fit/ener.py
@@ -182,6 +182,8 @@ def __init__(
                 self.atom_ener.append(None)
         self.useBN = False
         self.bias_atom_e = np.zeros(self.ntypes, dtype=np.float64)
+        ntypes_atom = self.ntypes - self.ntypes_spin
+        self.bias_atom_e = self.bias_atom_e[:ntypes_atom]
         self.register_buffer(
             "t_bias_atom_e",
             paddle.to_tensor(self.bias_atom_e),
@@ -259,7 +261,6 @@ def __init__(
                     1,
                     activation_fn=None,
                     precision=self.fitting_precision,
-                    bavg=self.bias_atom_e,
                     name=layer_suffix,
                     seed=self.seed,
                     trainable=self.trainable[-1],
@@ -321,6 +322,26 @@ def compute_output_stats(self, all_stat: dict, mixed_type: bool = False) -> None
         self.bias_atom_e = self._compute_output_stats(
             all_stat, rcond=self.rcond, mixed_type=mixed_type
         )
+        ntypes_atom = self.ntypes - self.ntypes_spin
+        if self.spin is not None:
+            for type_i in range(ntypes_atom):
+                if self.bias_atom_e.shape[0] != self.ntypes:
+                    self.bias_atom_e = np.pad(
+                        self.bias_atom_e,
+                        (0, self.ntypes_spin),
+                        "constant",
+                        constant_values=(0, 0),
+                    )
+                    bias_atom_e = self.bias_atom_e
+                if self.spin.use_spin[type_i]:
+                    self.bias_atom_e[type_i] = (
+                        self.bias_atom_e[type_i]
+                        + self.bias_atom_e[type_i + ntypes_atom]
+                    )
+                else:
+                    self.bias_atom_e[type_i] = self.bias_atom_e[type_i]
+            self.bias_atom_e = self.bias_atom_e[:ntypes_atom]
+
         paddle.assign(self.bias_atom_e, self.t_bias_atom_e)
 
     def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False):
@@ -525,26 +546,10 @@ def forward(
                 self.aparam_inv_std = 1.0
 
         ntypes_atom = self.ntypes - self.ntypes_spin
-        if self.spin is not None:
-            for type_i in range(ntypes_atom):
-                if self.bias_atom_e.shape[0] != self.ntypes:
-                    self.bias_atom_e = np.pad(
-                        self.bias_atom_e,
-                        (0, self.ntypes_spin),
-                        "constant",
-                        constant_values=(0, 0),
-                    )
-                    bias_atom_e = self.bias_atom_e
-                if self.spin.use_spin[type_i]:
-                    self.bias_atom_e[type_i] = (
-                        self.bias_atom_e[type_i]
-                        + self.bias_atom_e[type_i + ntypes_atom]
-                    )
-                else:
-                    self.bias_atom_e[type_i] = self.bias_atom_e[type_i]
-            self.bias_atom_e = self.bias_atom_e[:ntypes_atom]
 
-        inputs = paddle.reshape(inputs, [-1, natoms[0], self.dim_descrpt])
+        inputs = paddle.reshape(
+            inputs, [-1, natoms[0], self.dim_descrpt]
+        )  # [1, all_atoms, M1*M2]
         if len(self.atom_ener):
             # only for atom_ener
             nframes = input_dict.get("nframes")
@@ -558,7 +563,7 @@ def forward(
                 inputs_zero = paddle.zeros_like(inputs, dtype=GLOBAL_PD_FLOAT_PRECISION)
 
         if bias_atom_e is not None:
-            assert len(bias_atom_e) == self.ntypes
+            assert len(bias_atom_e) == self.ntypes - self.ntypes_spin
 
         fparam = None
         if self.numb_fparam > 0:
@@ -590,7 +595,7 @@ def forward(
                 atype_nall,
                 [0, 1],
                 [0, 0],
-                [-1, paddle.sum(natoms[2 : 2 + ntypes_atom]).item()],
+                [atype_nall.shape[0], paddle.sum(natoms[2 : 2 + ntypes_atom]).item()],
             )
             atype_filter = paddle.cast(self.atype_nloc >= 0, GLOBAL_PD_FLOAT_PRECISION)
             self.atype_nloc = paddle.reshape(self.atype_nloc, [-1])
diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py
index 27c6af754f..38b5401047 100644
--- a/deepmd/infer/deep_eval.py
+++ b/deepmd/infer/deep_eval.py
@@ -10,19 +10,20 @@
 
 import numpy as np
 
+import deepmd
 from deepmd.common import (
     j_loader,
     j_must_have,
 )
-from deepmd.descriptor import (
-    DescrptSeA,
-)
 from deepmd.env import (
     MODEL_VERSION,
     default_tf_session_config,
     paddle,
     tf,
 )
+from deepmd.fit import (
+    ener,
+)
 from deepmd.model import (
     EnerModel,
 )
@@ -32,6 +33,9 @@
 from deepmd.utils.sess import (
     run_sess,
 )
+from deepmd.utils.spin import (
+    Spin,
+)
 
 if TYPE_CHECKING:
     from pathlib import (
@@ -39,6 +43,30 @@
     )
 
 
+def remove_comment_in_json(jdata):
+    """Remove the comment in json file.
+
+    Parameters
+    ----------
+    jdata : dict
+        The data loaded from json file.
+
+    Returns
+    -------
+    dict
+        The new data without comments.
+    """
+    if not isinstance(jdata, dict):
+        return
+    find = False
+    for k, v in jdata.items():
+        if "_comment" == k:
+            find = True
+        remove_comment_in_json(v)
+    if find:
+        del jdata["_comment"]
+
+
 class DeepEval:
     """Common methods for DeepPot, DeepWFC, DeepPolar, ...
 
@@ -65,63 +93,224 @@ def __init__(
         auto_batch_size: Union[bool, int, AutoBatchSize] = False,
     ):
         jdata = j_loader("input.json")
+        remove_comment_in_json(jdata)
         model_param = j_must_have(jdata, "model")
-
+        self.multi_task_mode = "fitting_net_dict" in model_param
         descrpt_param = j_must_have(model_param, "descriptor")
+        fitting_param = (
+            j_must_have(model_param, "fitting_net")
+            if not self.multi_task_mode
+            else j_must_have(model_param, "fitting_net_dict")
+        )
+        typeebd_param = model_param.get("type_embedding", None)
+        spin_param = model_param.get("spin", None)
+        self.model_param = model_param
+        self.descrpt_param = descrpt_param
+        # spin
+        if spin_param is not None:
+            spin = Spin(
+                use_spin=spin_param["use_spin"],
+                virtual_len=spin_param["virtual_len"],
+                spin_norm=spin_param["spin_norm"],
+            )
+        else:
+            spin = None
+
+        # descriptor
+        try:
+            descrpt_type = descrpt_param["type"]
+            self.descrpt_type = descrpt_type
+        except KeyError:
+            raise KeyError("the type of descriptor should be set by `type`")
+
         explicit_ntypes_descrpt = ["se_atten"]
-        # hybrid_with_tebd = False
+        hybrid_with_tebd = False
         if descrpt_param["type"] in explicit_ntypes_descrpt:
             descrpt_param["ntypes"] = len(model_param["type_map"])
         elif descrpt_param["type"] == "hybrid":
             for descrpt_item in descrpt_param["list"]:
                 if descrpt_item["type"] in explicit_ntypes_descrpt:
                     descrpt_item["ntypes"] = len(model_param["type_map"])
-                    # hybrid_with_tebd = True
-
-        # if descrpt_param["type"] in ["se_e2_a", "se_a", "se_e2_r", "se_r", "hybrid"]:
-        descrpt_param["spin"] = None
-        descrpt_param["type_one_side"] = False
-
-        descrpt_param.pop("type", None)
-        descrpt_param.pop("_comment", None)
-        self.spin = None
-        # descrpt_param["spin"] = self.spin
-        self.descrpt = DescrptSeA(**descrpt_param)
-
-        self.multi_task_mode = "fitting_net_dict" in model_param
-        fitting_param = (
-            j_must_have(model_param, "fitting_net")
-            if not self.multi_task_mode
-            else j_must_have(model_param, "fitting_net_dict")
-        )
-        from deepmd.fit import (
-            EnerFitting,
-        )
-
-        # fitting_param.pop("type", None)
-        fitting_param.pop("_comment", None)
-        fitting_param["descrpt"] = self.descrpt
-        self.fitting = EnerFitting(**fitting_param)
-
-        self.typeebd = None
-
-        self.model = EnerModel(
-            self.descrpt,
-            self.fitting,
-            self.typeebd,
-            model_param.get("type_map"),
-            model_param.get("data_stat_nbatch", 10),
-            model_param.get("data_stat_protect", 1e-2),
-            model_param.get("use_srtab"),
-            model_param.get("smin_alpha"),
-            model_param.get("sw_rmin"),
-            model_param.get("sw_rmax"),
-            self.spin,
-        )
+                    hybrid_with_tebd = True
+        if self.multi_task_mode:
+            descrpt_param["multi_task"] = True
+        if descrpt_param["type"] in ["se_e2_a", "se_a", "se_e2_r", "se_r", "hybrid"]:
+            descrpt_param["spin"] = spin
+        descrpt_param.pop("type")
+        descrpt = deepmd.descriptor.se_a.DescrptSeA(**descrpt_param)
+
+        # fitting net
+        if not self.multi_task_mode:
+            fitting_type = fitting_param.get("type", "ener")
+            self.fitting_type = fitting_type
+            fitting_param["descrpt"] = descrpt
+            if fitting_type == "ener":
+                fitting_param["spin"] = spin
+                fitting_param.pop("type", None)
+            fitting = ener.EnerFitting(**fitting_param)
+        else:
+            self.fitting_dict = {}
+            self.fitting_type_dict = {}
+            self.nfitting = len(fitting_param)
+            for item in fitting_param:
+                item_fitting_param = fitting_param[item]
+                item_fitting_type = item_fitting_param.get("type", "ener")
+                self.fitting_type_dict[item] = item_fitting_type
+                item_fitting_param["descrpt"] = descrpt
+                if item_fitting_type == "ener":
+                    item_fitting_param["spin"] = spin
+                # self.fitting_dict[item] = Fitting(**item_fitting_param)
+
+        # type embedding
+        padding = False
+        if descrpt_type == "se_atten" or hybrid_with_tebd:
+            padding = True
+        if typeebd_param is not None:
+            raise NotImplementedError()
+            # typeebd = TypeEmbedNet(
+            #     neuron=typeebd_param["neuron"],
+            #     resnet_dt=typeebd_param["resnet_dt"],
+            #     activation_function=typeebd_param["activation_function"],
+            #     precision=typeebd_param["precision"],
+            #     trainable=typeebd_param["trainable"],
+            #     seed=typeebd_param["seed"],
+            #     padding=padding,
+            # )
+        elif descrpt_type == "se_atten" or hybrid_with_tebd:
+            raise NotImplementedError()
+            # default_args = type_embedding_args()
+            # default_args_dict = {i.name: i.default for i in default_args}
+            # typeebd = TypeEmbedNet(
+            #     neuron=default_args_dict["neuron"],
+            #     resnet_dt=default_args_dict["resnet_dt"],
+            #     activation_function=None,
+            #     precision=default_args_dict["precision"],
+            #     trainable=default_args_dict["trainable"],
+            #     seed=default_args_dict["seed"],
+            #     padding=padding,
+            # )
+        else:
+            typeebd = None
+
+        # init model
+        # infer model type by fitting_type
+        if not self.multi_task_mode:
+            if self.fitting_type == "ener":
+                self.model = EnerModel(
+                    descrpt,
+                    fitting,
+                    typeebd,
+                    model_param.get("type_map"),
+                    model_param.get("data_stat_nbatch", 10),
+                    model_param.get("data_stat_protect", 1e-2),
+                    model_param.get("use_srtab"),
+                    model_param.get("smin_alpha"),
+                    model_param.get("sw_rmin"),
+                    model_param.get("sw_rmax"),
+                    spin,
+                )
+            # elif fitting_type == 'wfc':
+            #     self.model = WFCModel(model_param, descrpt, fitting)
+            elif self.fitting_type == "dos":
+                raise NotImplementedError()
+                # self.model = DOSModel(
+                #     descrpt,
+                #     fitting,
+                #     typeebd,
+                #     model_param.get("type_map"),
+                #     model_param.get("data_stat_nbatch", 10),
+                #     model_param.get("data_stat_protect", 1e-2),
+                # )
+
+            elif self.fitting_type == "dipole":
+                raise NotImplementedError()
+                # self.model = DipoleModel(
+                #     descrpt,
+                #     fitting,
+                #     typeebd,
+                #     model_param.get("type_map"),
+                #     model_param.get("data_stat_nbatch", 10),
+                #     model_param.get("data_stat_protect", 1e-2),
+                # )
+            elif self.fitting_type == "polar":
+                raise NotImplementedError()
+                # self.model = PolarModel(
+                #     descrpt,
+                #     fitting,
+                #     typeebd,
+                #     model_param.get("type_map"),
+                #     model_param.get("data_stat_nbatch", 10),
+                #     model_param.get("data_stat_protect", 1e-2),
+                # )
+            # elif self.fitting_type == 'global_polar':
+            #     self.model = GlobalPolarModel(
+            #         descrpt,
+            #         fitting,
+            #         model_param.get('type_map'),
+            #         model_param.get('data_stat_nbatch', 10),
+            #         model_param.get('data_stat_protect', 1e-2)
+            #     )
+            else:
+                raise RuntimeError("get unknown fitting type when building model")
+        else:  # multi-task mode
+            raise NotImplementedError()
+            # self.model = MultiModel(
+            #     descrpt,
+            #     self.fitting_dict,
+            #     self.fitting_type_dict,
+            #     typeebd,
+            #     model_param.get("type_map"),
+            #     model_param.get("data_stat_nbatch", 10),
+            #     model_param.get("data_stat_protect", 1e-2),
+            #     model_param.get("use_srtab"),
+            #     model_param.get("smin_alpha"),
+            #     model_param.get("sw_rmin"),
+            #     model_param.get("sw_rmax"),
+            # )
+
+        # # if descrpt_param["type"] in ["se_e2_a", "se_a", "se_e2_r", "se_r", "hybrid"]:
+        # descrpt_param["spin"] = None
+        # descrpt_param["type_one_side"] = False
+
+        # descrpt_param.pop("type", None)
+        # descrpt_param.pop("_comment", None)
+        # spin = None
+        # # descrpt_param["spin"] = spin
+        # descrpt = DescrptSeA(**descrpt_param)
+
+        # self.multi_task_mode = "fitting_net_dict" in model_param
+        # fitting_param = (
+        #     j_must_have(model_param, "fitting_net")
+        #     if not self.multi_task_mode
+        #     else j_must_have(model_param, "fitting_net_dict")
+        # )
+        # from deepmd.fit import EnerFitting
+
+        # # fitting_param.pop("type", None)
+        # fitting_param.pop("_comment", None)
+        # fitting_param["descrpt"] = descrpt
+        # fitting = EnerFitting(**fitting_param)
+
+        # typeebd = None
+
+        # self.model = EnerModel(
+        #     descrpt,
+        #     fitting,
+        #     typeebd,
+        #     model_param.get("type_map"),
+        #     model_param.get("data_stat_nbatch", 10),
+        #     model_param.get("data_stat_protect", 1e-2),
+        #     model_param.get("use_srtab"),
+        #     model_param.get("smin_alpha"),
+        #     model_param.get("sw_rmin"),
+        #     model_param.get("sw_rmax"),
+        #     spin,
+        # )
         model_file_str = str(model_file)
         if model_file_str.endswith((".pdmodel", ".pdiparams")):
             st_model_prefix = model_file_str.rsplit(".", 1)[0]
             self.st_model = paddle.jit.load(st_model_prefix)
+            print(f"==>> Load static model successfully from: {str(st_model_prefix)}")
         else:
             load_state_dict = paddle.load(str(model_file))
             for k, v in load_state_dict.items():
@@ -143,7 +332,7 @@ def __init__(
                             self.model.state_dict()[k].shape
                         )
             self.model.set_state_dict(load_state_dict)
-        print(f"==>> Load pretraied model successfully from: {str(model_file)}")
+            print(f"==>> Load dynamic model successfully from: {str(model_file)}")
         self.load_prefix = load_prefix
 
         # graph_compatable should be called after graph and prefix are set
diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py
index 377c776320..11a7ab356e 100644
--- a/deepmd/infer/deep_pot.py
+++ b/deepmd/infer/deep_pot.py
@@ -113,11 +113,12 @@ def __init__(
                 if attr_name != "t_descriptor":
                     raise
 
-        self.ntypes = 2
-        self.rcut = 6.0
+        self.ntypes = int(self.model.descrpt.buffer_ntypes)
+        self.rcut = float(self.model.descrpt.buffer_rcut)
         self.dfparam = 0
         self.daparam = 0
-        self.t_tmap = ["O", "H"]
+        self.t_tmap = [chr(idx) for idx in self.model.buffer_tmap.tolist()]
+        self.t_tmap = [c for c in self.t_tmap if c != " "]
 
         # setup modifier
         try:
diff --git a/deepmd/loss/ener.py b/deepmd/loss/ener.py
index d11177ee3a..8427f4ec6a 100644
--- a/deepmd/loss/ener.py
+++ b/deepmd/loss/ener.py
@@ -216,8 +216,6 @@ def compute_loss(self, learning_rate, natoms, model_dict, label_dict, suffix):
         return l2_loss, more_loss
 
     def eval(self, model, batch_data, natoms):
-        # placeholder = self.l2_l
-
         model_inputs = {}
         for kk in batch_data.keys():
             if kk == "find_type" or kk == "type":
@@ -340,7 +338,7 @@ def __init__(
                 default=1.0,
             )
 
-    def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
+    def compute_loss(self, learning_rate, natoms, model_dict, label_dict, suffix):
         energy_pred = model_dict["energy"]
         force_pred = model_dict["force"]
         virial_pred = model_dict["virial"]
@@ -365,151 +363,202 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix):
             # E = - E(A) - E(B) + E(C) + E(D)
             # A, B, C, D could be put far away from each other
             atom_ener_coeff = label_dict["atom_ener_coeff"]
-            atom_ener_coeff = tf.reshape(atom_ener_coeff, tf.shape(atom_ener_pred))
-            energy_pred = tf.reduce_sum(atom_ener_coeff * atom_ener_pred, 1)
-        l2_ener_loss = tf.reduce_mean(
-            tf.square(energy_pred - energy_label), name="l2_" + suffix
+            atom_ener_coeff = paddle.reshape(
+                atom_ener_coeff, paddle.shape(atom_ener_pred)
+            )
+            energy_pred = paddle.sum(atom_ener_coeff * atom_ener_pred, 1)
+        l2_ener_loss = (
+            paddle.mean(paddle.square(energy_pred - energy_label), name="l2_" + suffix)
+            if self.has_e
+            else 0.0
         )
 
         # split force to force_r and force_m, compute their respective loss
-        real_nloc = tf.reduce_sum(natoms[2 : 2 + len(self.use_spin)])
+        real_nloc = paddle.sum(natoms[2 : 2 + len(self.use_spin)])
         virt_nloc = natoms[0] - real_nloc
-        force_pred_reshape = tf.reshape(force_pred, [-1, natoms[0] * 3])
-        force_label_reshape = tf.reshape(force_label, [-1, natoms[0] * 3])
-        force_r_pred = tf.reshape(
-            tf.slice(force_pred_reshape, [0, 0], [-1, real_nloc * 3]), [-1]
-        )
-        force_m_pred = tf.reshape(
-            tf.slice(force_pred_reshape, [0, real_nloc * 3], [-1, virt_nloc * 3]), [-1]
-        )
-        force_r_label = tf.reshape(
-            tf.slice(force_label_reshape, [0, 0], [-1, real_nloc * 3]), [-1]
-        )
-        force_m_label = tf.reshape(
-            tf.slice(force_label_reshape, [0, real_nloc * 3], [-1, virt_nloc * 3]), [-1]
+        force_pred_reshape = paddle.reshape(force_pred, [-1, natoms[0] * 3])
+        force_label_reshape = paddle.reshape(force_label, [-1, natoms[0] * 3])
+        force_r_pred = paddle.reshape(
+            paddle.slice(
+                force_pred_reshape,
+                [0, 1],
+                [0, 0],
+                [force_pred_reshape.shape[0], real_nloc * 3],
+            ),
+            [-1],
+        )
+        force_m_pred = paddle.reshape(
+            paddle.slice(
+                force_pred_reshape,
+                [0, 1],
+                [0, real_nloc * 3],
+                [force_pred_reshape.shape[0], real_nloc * 3 + virt_nloc * 3],
+            ),
+            [-1],
+        )
+        force_r_label = paddle.reshape(
+            paddle.slice(
+                force_label_reshape,
+                [0, 1],
+                [0, 0],
+                [force_label_reshape.shape[0], real_nloc * 3],
+            ),
+            [-1],
+        )
+        force_m_label = paddle.reshape(
+            paddle.slice(
+                force_label_reshape,
+                [0, 1],
+                [0, real_nloc * 3],
+                [force_label_reshape.shape[0], real_nloc * 3 + virt_nloc * 3],
+            ),
+            [-1],
         )
-        l2_force_r_loss = tf.reduce_mean(
-            tf.square(force_r_pred - force_r_label), name="l2_force_real_" + suffix
+        l2_force_r_loss = (
+            paddle.mean(
+                paddle.square(force_r_pred - force_r_label),
+                name="l2_force_real_" + suffix,
+            )
+            if self.has_fr
+            else 0.0
         )
-        l2_force_m_loss = tf.reduce_mean(
-            tf.square(force_m_pred - force_m_label), name="l2_force_mag_" + suffix
+        l2_force_m_loss = (
+            paddle.mean(
+                paddle.square(force_m_pred - force_m_label),
+                name="l2_force_mag_" + suffix,
+            )
+            if self.has_fm
+            else 0.0
         )
 
-        virial_pred_reshape = tf.reshape(virial_pred, [-1])
-        virial_label_reshape = tf.reshape(virial_label, [-1])
-        l2_virial_loss = tf.reduce_mean(
-            tf.square(virial_pred_reshape - virial_label_reshape),
-            name="l2_virial_" + suffix,
+        virial_pred_reshape = paddle.reshape(virial_pred, [-1])
+        virial_label_reshape = paddle.reshape(virial_label, [-1])
+        l2_virial_loss = (
+            paddle.mean(
+                paddle.square(virial_pred_reshape - virial_label_reshape),
+                name="l2_virial_" + suffix,
+            )
+            if self.has_v
+            else 0.0
         )
 
         # need to change. can't get atom_ener_hat?
-        atom_ener_pred_reshape = tf.reshape(atom_ener_pred, [-1])
-        atom_ener_label_reshape = tf.reshape(atom_ener_label, [-1])
-        l2_atom_ener_loss = tf.reduce_mean(
-            tf.square(atom_ener_pred_reshape - atom_ener_label_reshape),
-            name="l2_atom_ener_" + suffix,
+        atom_ener_pred_reshape = paddle.reshape(atom_ener_pred, [-1])
+        atom_ener_label_reshape = paddle.reshape(atom_ener_label, [-1])
+        l2_atom_ener_loss = (
+            paddle.mean(
+                paddle.square(atom_ener_pred_reshape - atom_ener_label_reshape),
+                name="l2_atom_ener_" + suffix,
+            )
+            if self.has_ae
+            else 0.0
         )
 
-        atom_norm = 1.0 / global_cvt_2_tf_float(natoms[0])
-        atom_norm_ener = 1.0 / global_cvt_2_ener_float(natoms[0])
-        pref_e = global_cvt_2_ener_float(
-            find_energy
-            * (
-                self.limit_pref_e
-                + (self.start_pref_e - self.limit_pref_e)
-                * learning_rate
-                / self.starter_learning_rate
-            )
+        atom_norm = 1.0 / (natoms[0])
+        atom_norm_ener = 1.0 / (natoms[0])
+        pref_e = find_energy * (
+            self.limit_pref_e
+            + (self.start_pref_e - self.limit_pref_e)
+            * learning_rate
+            / self.starter_learning_rate
         )
-        pref_fr = global_cvt_2_tf_float(
-            find_force
-            * (
-                self.limit_pref_fr
-                + (self.start_pref_fr - self.limit_pref_fr)
-                * learning_rate
-                / self.starter_learning_rate
-            )
+        pref_fr = find_force * (
+            self.limit_pref_fr
+            + (self.start_pref_fr - self.limit_pref_fr)
+            * learning_rate
+            / self.starter_learning_rate
         )
-        pref_fm = global_cvt_2_tf_float(
-            find_force
-            * (
-                self.limit_pref_fm
-                + (self.start_pref_fm - self.limit_pref_fm)
-                * learning_rate
-                / self.starter_learning_rate
-            )
+        pref_fm = find_force * (
+            self.limit_pref_fm
+            + (self.start_pref_fm - self.limit_pref_fm)
+            * learning_rate
+            / self.starter_learning_rate
         )
-        pref_v = global_cvt_2_tf_float(
-            find_virial
-            * (
-                self.limit_pref_v
-                + (self.start_pref_v - self.limit_pref_v)
-                * learning_rate
-                / self.starter_learning_rate
-            )
+        pref_v = find_virial * (
+            self.limit_pref_v
+            + (self.start_pref_v - self.limit_pref_v)
+            * learning_rate
+            / self.starter_learning_rate
         )
-        pref_ae = global_cvt_2_tf_float(
-            find_atom_ener
-            * (
-                self.limit_pref_ae
-                + (self.start_pref_ae - self.limit_pref_ae)
-                * learning_rate
-                / self.starter_learning_rate
-            )
+        pref_ae = find_atom_ener * (
+            self.limit_pref_ae
+            + (self.start_pref_ae - self.limit_pref_ae)
+            * learning_rate
+            / self.starter_learning_rate
         )
 
         l2_loss = 0
         more_loss = {}
         if self.has_e:
             l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
-        more_loss["l2_ener_loss"] = l2_ener_loss
+            more_loss["l2_ener_loss"] = l2_ener_loss
         if self.has_fr:
-            l2_loss += global_cvt_2_ener_float(pref_fr * l2_force_r_loss)
-        more_loss["l2_force_r_loss"] = l2_force_r_loss
+            l2_loss += pref_fr * l2_force_r_loss
+            more_loss["l2_force_r_loss"] = l2_force_r_loss
         if self.has_fm:
-            l2_loss += global_cvt_2_ener_float(pref_fm * l2_force_m_loss)
-        more_loss["l2_force_m_loss"] = l2_force_m_loss
+            l2_loss += pref_fm * l2_force_m_loss
+            more_loss["l2_force_m_loss"] = l2_force_m_loss
         if self.has_v:
-            l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss))
-        more_loss["l2_virial_loss"] = l2_virial_loss
+            l2_loss += atom_norm * (pref_v * l2_virial_loss)
+            more_loss["l2_virial_loss"] = l2_virial_loss
         if self.has_ae:
-            l2_loss += global_cvt_2_ener_float(pref_ae * l2_atom_ener_loss)
-        more_loss["l2_atom_ener_loss"] = l2_atom_ener_loss
-
-        # only used when tensorboard was set as true
-        self.l2_loss_summary = tf.summary.scalar("l2_loss", tf.sqrt(l2_loss))
-        self.l2_loss_ener_summary = tf.summary.scalar(
-            "l2_ener_loss",
-            global_cvt_2_tf_float(tf.sqrt(l2_ener_loss))
-            / global_cvt_2_tf_float(natoms[0]),
-        )
-        self.l2_loss_force_r_summary = tf.summary.scalar(
-            "l2_force_r_loss", tf.sqrt(l2_force_r_loss)
-        )
-        self.l2_loss_force_m_summary = tf.summary.scalar(
-            "l2_force_m_loss", tf.sqrt(l2_force_m_loss)
-        )
-        self.l2_loss_virial_summary = tf.summary.scalar(
-            "l2_virial_loss", tf.sqrt(l2_virial_loss) / global_cvt_2_tf_float(natoms[0])
-        )
+            l2_loss += pref_ae * l2_atom_ener_loss
+            more_loss["l2_atom_ener_loss"] = l2_atom_ener_loss
 
         self.l2_l = l2_loss
         self.l2_more = more_loss
         return l2_loss, more_loss
 
-    def eval(self, sess, feed_dict, natoms):
-        placeholder = self.l2_l
+    def eval(self, model, batch_data, natoms):
+        model_inputs = {}
+        for kk in batch_data.keys():
+            if kk == "find_type" or kk == "type":
+                continue
+            prec = "float64"
+            if "find_" in kk:
+                model_inputs[kk] = paddle.to_tensor(batch_data[kk], dtype="float64")
+            else:
+                model_inputs[kk] = paddle.to_tensor(
+                    np.reshape(batch_data[kk], [-1]), dtype=prec
+                )
+
+        for ii in ["type"]:
+            model_inputs[ii] = paddle.to_tensor(
+                np.reshape(batch_data[ii], [-1]), dtype="int32"
+            )
+        for ii in ["natoms_vec", "default_mesh"]:
+            model_inputs[ii] = paddle.to_tensor(batch_data[ii], dtype="int32")
+        model_inputs["is_training"] = paddle.to_tensor(False)
+        model_inputs["natoms_vec"] = paddle.to_tensor(
+            model_inputs["natoms_vec"], place="cpu"
+        )
+
+        model_pred = model(
+            model_inputs["coord"],
+            model_inputs["type"],
+            model_inputs["natoms_vec"],
+            model_inputs["box"],
+            model_inputs["default_mesh"],
+            model_inputs,
+            suffix="",
+            reuse=False,
+        )
+        l2_l, l2_more = self.compute_loss(
+            0.0,
+            model_inputs["natoms_vec"],
+            model_pred,
+            model_inputs,
+            suffix="test",
+        )
         run_data = [
-            self.l2_l,
-            self.l2_more["l2_ener_loss"] if self.has_e else placeholder,
-            self.l2_more["l2_force_r_loss"] if self.has_fr else placeholder,
-            self.l2_more["l2_force_m_loss"] if self.has_fm else placeholder,
-            self.l2_more["l2_virial_loss"] if self.has_v else placeholder,
-            self.l2_more["l2_atom_ener_loss"] if self.has_ae else placeholder,
+            (float(l2_l)),
+            (float(l2_more["l2_ener_loss"]) if self.has_e else 0.0),
+            (float(l2_more["l2_force_r_loss"]) if self.has_fr else 0.0),
+            (float(l2_more["l2_force_m_loss"]) if self.has_fm else 0.0),
+            (float(l2_more["l2_virial_loss"]) if self.has_v else 0.0),
+            (float(l2_more["l2_atom_ener_loss"]) if self.has_ae else 0.0),
         ]
-        error, error_e, error_fr, error_fm, error_v, error_ae = run_sess(
-            sess, run_data, feed_dict=feed_dict
-        )
+        error, error_e, error_fr, error_fm, error_v, error_ae = run_data
         results = {"natoms": natoms[0], "rmse": np.sqrt(error)}
         if self.has_e:
             results["rmse_e"] = np.sqrt(error_e) / natoms[0]
diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py
index 21c5ec6ee0..d2450b5a8d 100644
--- a/deepmd/model/ener.py
+++ b/deepmd/model/ener.py
@@ -4,6 +4,8 @@
     Optional,
 )
 
+import numpy as np
+
 from deepmd.env import (
     MODEL_VERSION,
     paddle,
@@ -25,9 +27,9 @@
 )
 
 if TYPE_CHECKING:
-    from deepmd.fit import (
-        ener,
-    )
+    import paddle  # noqa: F811
+
+    from deepmd.fit import ener  # noqa: F811
 
 
 class EnerModel(Model, paddle.nn.Layer):
@@ -106,16 +108,27 @@ def __init__(
         self.t_ver = str(MODEL_VERSION)
         # NOTE: workaround for string type is not supported in Paddle
         self.register_buffer(
-            "buffer_t_type",
+            "buffer_tmap",
             paddle.to_tensor([ord(c) for c in self.t_tmap], dtype="int32"),
         )
         self.register_buffer(
-            "buffer_t_mt", paddle.to_tensor([ord(c) for c in self.t_mt], dtype="int32")
+            "buffer_model_type",
+            paddle.to_tensor([ord(c) for c in self.t_mt], dtype="int32"),
         )
         self.register_buffer(
-            "buffer_t_ver",
+            "buffer_model_version",
             paddle.to_tensor([ord(c) for c in self.t_ver], dtype="int32"),
         )
+        if self.srtab is not None:
+            tab_info, tab_data = self.srtab.get()
+            self.tab_info = paddle.register_buffer(
+                "buffer_t_tab_info",
+                paddle.to_tensor(tab_info, dtype=paddle.float64),
+            )
+            self.tab_data = paddle.register_buffer(
+                "buffer_t_tab_data",
+                paddle.to_tensor(tab_data, dtype=paddle.float64),
+            )
 
     def get_rcut(self):
         return self.rcut
@@ -194,10 +207,52 @@ def forward(
             reuse=reuse,
         )
 
+        if self.srtab is not None:
+            nlist, rij, sel_a, sel_r = self.descrpt.get_nlist()
+            nnei_a = np.cumsum(sel_a)[-1]
+            nnei_r = np.cumsum(sel_r)[-1]
+
         atom_ener = self.fitting(dout, natoms, input_dict, reuse=reuse, suffix=suffix)
         self.atom_ener = atom_ener
 
-        energy_raw = atom_ener
+        if self.srtab is not None:
+            raise NotImplementedError(
+                f"srtab not implemented in {self.__class__.__name__}"
+            )
+            # sw_lambda, sw_deriv = op_module.soft_min_switch(
+            #     atype,
+            #     rij,
+            #     nlist,
+            #     natoms,
+            #     sel_a=sel_a,
+            #     sel_r=sel_r,
+            #     alpha=self.smin_alpha,
+            #     rmin=self.sw_rmin,
+            #     rmax=self.sw_rmax,
+            # )
+            # inv_sw_lambda = 1.0 - sw_lambda
+            # # NOTICE:
+            # # atom energy is not scaled,
+            # # force and virial are scaled
+            # tab_atom_ener, tab_force, tab_atom_virial = op_module.pair_tab(
+            #     self.tab_info,
+            #     self.tab_data,
+            #     atype,
+            #     rij,
+            #     nlist,
+            #     natoms,
+            #     sw_lambda,
+            #     sel_a=sel_a,
+            #     sel_r=sel_r,
+            # )
+            # energy_diff = tab_atom_ener - tf.reshape(atom_ener, [-1, natoms[0]])
+            # tab_atom_ener = tf.reshape(sw_lambda, [-1]) * tf.reshape(
+            #     tab_atom_ener, [-1]
+            # )
+            # atom_ener = tf.reshape(inv_sw_lambda, [-1]) * atom_ener
+            # energy_raw = tab_atom_ener + atom_ener
+        else:
+            energy_raw = atom_ener
 
         nloc_atom = (
             natoms[0]
@@ -214,18 +269,42 @@ def forward(
         # virial: [1, 9]
         # atom_virial: [1, all_atoms*9]
 
-        force = paddle.reshape(force, [-1, 3 * natoms[1]])
+        if self.srtab is not None:
+            raise NotImplementedError()
+            # sw_force = op_module.soft_min_force(
+            #     energy_diff, sw_deriv, nlist, natoms, n_a_sel=nnei_a, n_r_sel=nnei_r
+            # )
+            # force = force + sw_force + tab_force
+
+        force = paddle.reshape(force, [-1, 3 * natoms[1]])  # [1, all_atoms*3]
         if self.spin is not None:
             # split and concatenate force to compute local atom force and magnetic force
             judge = paddle.equal(natoms[0], natoms[1])
-            force = paddle.where(
-                judge,
-                self.natoms_match(force, natoms),
-                self.natoms_not_match(force, natoms, atype),
-            )
+            if judge.item():
+                force = self.natoms_match(force, natoms)
+            else:
+                force = self.natoms_not_match(force, natoms, atype)
 
         force = paddle.reshape(force, [-1, 3 * natoms[1]], name="o_force" + suffix)
 
+        if self.srtab is not None:
+            raise NotImplementedError()
+            # sw_virial, sw_atom_virial = op_module.soft_min_virial(
+            #     energy_diff,
+            #     sw_deriv,
+            #     rij,
+            #     nlist,
+            #     natoms,
+            #     n_a_sel=nnei_a,
+            #     n_r_sel=nnei_r,
+            # )
+            # atom_virial = atom_virial + sw_atom_virial + tab_atom_virial
+            # virial = (
+            #     virial
+            #     + sw_virial
+            #     + tf.sum(tf.reshape(tab_atom_virial, [-1, natoms[1], 9]), axis=1)
+            # )
+
         virial = paddle.reshape(virial, [-1, 9], name="o_virial" + suffix)
         atom_virial = paddle.reshape(
             atom_virial, [-1, 9 * natoms[1]], name="o_atom_virial" + suffix
@@ -279,41 +358,69 @@ def natoms_match(self, force, natoms):
         use_spin = self.spin.use_spin
         virtual_len = self.spin.virtual_len
         spin_norm = self.spin.spin_norm
-        natoms_index = tf.concat([[0], tf.cumsum(natoms[2:])], axis=0)
+        natoms_index = paddle.concat(
+            [
+                paddle.to_tensor([0], dtype=natoms.dtype),
+                paddle.cumsum(natoms[2:]),
+            ],
+            axis=0,
+        )
         force_real_list = []
         for idx, use in enumerate(use_spin):
             if use is True:
                 force_real_list.append(
-                    tf.slice(
-                        force, [0, natoms_index[idx] * 3], [-1, natoms[idx + 2] * 3]
+                    paddle.slice(
+                        force,
+                        [0, 1],
+                        [0, natoms_index[idx] * 3],
+                        [
+                            force.shape[0],
+                            natoms_index[idx] * 3 + natoms[idx + 2].item() * 3,
+                        ],
                     )
-                    + tf.slice(
+                    + paddle.slice(
                         force,
+                        [0, 1],
                         [0, natoms_index[idx + len(use_spin)] * 3],
-                        [-1, natoms[idx + 2 + len(use_spin)] * 3],
+                        [
+                            force.shape[0],
+                            natoms_index[idx + len(use_spin)] * 3
+                            + natoms[idx + 2 + len(use_spin)].item() * 3,
+                        ],
                     )
                 )
             else:
                 force_real_list.append(
-                    tf.slice(
-                        force, [0, natoms_index[idx] * 3], [-1, natoms[idx + 2] * 3]
+                    paddle.slice(
+                        force,
+                        [0, 1],
+                        [0, natoms_index[idx] * 3],
+                        [
+                            force.shape[0],
+                            natoms_index[idx] * 3 + natoms[idx + 2].item() * 3,
+                        ],
                     )
                 )
         force_mag_list = []
         for idx, use in enumerate(use_spin):
             if use is True:
                 force_mag_list.append(
-                    tf.slice(
+                    paddle.slice(
                         force,
+                        [0, 1],
                         [0, natoms_index[idx + len(use_spin)] * 3],
-                        [-1, natoms[idx + 2 + len(use_spin)] * 3],
+                        [
+                            force.shape[0],
+                            natoms_index[idx + len(use_spin)] * 3
+                            + natoms[idx + 2 + len(use_spin)].item() * 3,
+                        ],
                     )
                 )
                 force_mag_list[idx] *= virtual_len[idx] / spin_norm[idx]
 
-        force_real = tf.concat(force_real_list, axis=1)
-        force_mag = tf.concat(force_mag_list, axis=1)
-        loc_force = tf.concat([force_real, force_mag], axis=1)
+        force_real = paddle.concat(force_real_list, axis=1)
+        force_mag = paddle.concat(force_mag_list, axis=1)
+        loc_force = paddle.concat([force_real, force_mag], axis=1)
         force = loc_force
         return force
 
@@ -325,48 +432,82 @@ def natoms_not_match(self, force, natoms, atype):
         spin_norm = self.spin.spin_norm
         loc_force = self.natoms_match(force, natoms)
         aatype = atype[0, :]
+
+        # FIXME: paddle.unique和tf.unique的返回值顺序是不一致, 下面的代码实现不正确，
+        # 在某些案例中需要修改
         ghost_atype = aatype[natoms[0] :]
-        _, _, ghost_natoms = tf.unique_with_counts(ghost_atype)
-        ghost_natoms_index = tf.concat([[0], tf.cumsum(ghost_natoms)], axis=0)
+        _, idx, ghost_natoms = paddle.unique(
+            ghost_atype, return_index=True, return_counts=True
+        )
+        idx_inv = paddle.empty_like(idx)
+        # NOTE: Use inverse permutation to get equaivalent result as tf
+        idx_inv[idx] = paddle.arange(0, len(idx))
+        ghost_natoms = ghost_natoms[idx_inv]
+
+        ghost_natoms_index = paddle.concat(
+            [
+                paddle.to_tensor([0], dtype=ghost_natoms.dtype),
+                paddle.cumsum(ghost_natoms),
+            ],
+            axis=0,
+        )
         ghost_natoms_index += natoms[0]
 
         ghost_force_real_list = []
         for idx, use in enumerate(use_spin):
             if use is True:
                 ghost_force_real_list.append(
-                    tf.slice(
+                    paddle.slice(
                         force,
+                        [0, 1],
                         [0, ghost_natoms_index[idx] * 3],
-                        [-1, ghost_natoms[idx] * 3],
+                        [
+                            force.shape[0],
+                            ghost_natoms_index[idx] * 3 + ghost_natoms[idx].item() * 3,
+                        ],
                     )
-                    + tf.slice(
+                    + paddle.slice(
                         force,
+                        [0, 1],
                         [0, ghost_natoms_index[idx + len(use_spin)] * 3],
-                        [-1, ghost_natoms[idx + len(use_spin)] * 3],
+                        [
+                            force.shape[0],
+                            ghost_natoms_index[idx + len(use_spin)] * 3
+                            + ghost_natoms[idx + len(use_spin)].item() * 3,
+                        ],
                     )
                 )
             else:
                 ghost_force_real_list.append(
-                    tf.slice(
+                    paddle.slice(
                         force,
+                        [0, 1],
                         [0, ghost_natoms_index[idx] * 3],
-                        [-1, ghost_natoms[idx] * 3],
+                        [
+                            force.shape[0],
+                            ghost_natoms_index[idx] * 3 + ghost_natoms[idx].item() * 3,
+                        ],
                     )
                 )
         ghost_force_mag_list = []
         for idx, use in enumerate(use_spin):
             if use is True:
                 ghost_force_mag_list.append(
-                    tf.slice(
+                    paddle.slice(
                         force,
+                        [0, 1],
                         [0, ghost_natoms_index[idx + len(use_spin)] * 3],
-                        [-1, ghost_natoms[idx + len(use_spin)] * 3],
+                        [
+                            force.shape[0],
+                            ghost_natoms_index[idx + len(use_spin)] * 3
+                            + ghost_natoms[idx + len(use_spin)].item() * 3,
+                        ],
                     )
                 )
                 ghost_force_mag_list[idx] *= virtual_len[idx] / spin_norm[idx]
 
-        ghost_force_real = tf.concat(ghost_force_real_list, axis=1)
-        ghost_force_mag = tf.concat(ghost_force_mag_list, axis=1)
-        ghost_force = tf.concat([ghost_force_real, ghost_force_mag], axis=1)
-        force = tf.concat([loc_force, ghost_force], axis=1)
+        ghost_force_real = paddle.concat(ghost_force_real_list, axis=1)
+        ghost_force_mag = paddle.concat(ghost_force_mag_list, axis=1)
+        ghost_force = paddle.concat([ghost_force_real, ghost_force_mag], axis=1)
+        force = paddle.concat([loc_force, ghost_force], axis=1)
         return force
diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py
index cc46dc9801..51232053e2 100644
--- a/deepmd/train/trainer.py
+++ b/deepmd/train/trainer.py
@@ -913,16 +913,7 @@ def train(self, train_data=None, valid_data=None, stop_batch: int = 10):
             # use tensorboard to visualize the training of deepmd-kit
             # it will takes some extra execution time to generate the tensorboard data
             if self.tensorboard and (cur_batch % self.tensorboard_freq == 0):
-                model_pred = self.model(
-                    paddle.to_tensor(train_batch["coord"], "float32"),
-                    paddle.to_tensor(train_batch["type"], "int32"),
-                    paddle.to_tensor(train_batch["natoms_vec"], "int32", "cpu"),
-                    paddle.to_tensor(train_batch["box"], "float32"),
-                    paddle.to_tensor(train_batch["default_mesh"], "int32"),
-                    train_batch,
-                    suffix="",
-                    reuse=False,
-                )
+                pass
             else:
                 model_inputs = {}
                 for kk in train_batch.keys():
@@ -958,7 +949,6 @@ def train(self, train_data=None, valid_data=None, stop_batch: int = 10):
                     suffix="",
                     reuse=False,
                 )
-
                 # print(f"{self.cur_batch} {self.learning_rate.get_lr():.10f}")
                 l2_l, l2_more = self.loss.compute_loss(
                     self.learning_rate.get_lr(),
diff --git a/deepmd/utils/neighbor_stat.py b/deepmd/utils/neighbor_stat.py
index 966645996a..ef26b47949 100644
--- a/deepmd/utils/neighbor_stat.py
+++ b/deepmd/utils/neighbor_stat.py
@@ -9,6 +9,7 @@
 import paddle
 
 from deepmd.env import (
+    GLOBAL_PD_FLOAT_PRECISION,
     op_module,
 )
 from deepmd.utils.data_system import (
@@ -71,25 +72,21 @@ def get_stat(self, data: DeepmdDataSystem) -> Tuple[float, List[int]]:
                     coord = np.array(data_set["coord"])[kk].reshape(
                         [-1, data.natoms[ii] * 3]
                     )
-                    coord = paddle.to_tensor(coord, dtype="float32", place="cpu")
+                    coord = paddle.to_tensor(coord, GLOBAL_PD_FLOAT_PRECISION, "cpu")
 
                     _type = np.array(data_set["type"])[kk].reshape(
                         [-1, data.natoms[ii]]
                     )
-                    _type = paddle.to_tensor(_type, dtype="int32", place="cpu")
+                    _type = paddle.to_tensor(_type, "int32", "cpu")
 
                     natoms_vec = np.array(data.natoms_vec[ii])
-                    natoms_vec = paddle.to_tensor(
-                        natoms_vec, dtype="int64", place="cpu"
-                    )
+                    natoms_vec = paddle.to_tensor(natoms_vec, "int32", "cpu")
 
                     box = np.array(data_set["box"])[kk].reshape([-1, 9])
-                    box = paddle.to_tensor(box, dtype="float32", place="cpu")
+                    box = paddle.to_tensor(box, GLOBAL_PD_FLOAT_PRECISION, "cpu")
 
                     default_mesh = np.array(data.default_mesh[ii])
-                    default_mesh = paddle.to_tensor(
-                        default_mesh, dtype="int32", place="cpu"
-                    )
+                    default_mesh = paddle.to_tensor(default_mesh, "int32", "cpu")
 
                     rcut = self.rcut
                     mn, dt = op_module.neighbor_stat(
diff --git a/deepmd/utils/network.py b/deepmd/utils/network.py
index 58e6378215..9515e1cad1 100644
--- a/deepmd/utils/network.py
+++ b/deepmd/utils/network.py
@@ -5,7 +5,6 @@
 )
 from deepmd.env import (
     GLOBAL_PD_FLOAT_PRECISION,
-    GLOBAL_TF_FLOAT_PRECISION,
     paddle,
     tf,
 )
@@ -19,7 +18,7 @@ def one_layer(
     inputs,
     outputs_size,
     activation_fn=tf.nn.tanh,
-    precision=GLOBAL_TF_FLOAT_PRECISION,
+    precision=GLOBAL_PD_FLOAT_PRECISION,
     stddev=1.0,
     bavg=0.0,
     name="linear",
@@ -321,7 +320,6 @@ def __init__(
         self.use_timestep = use_timestep
         self.useBN = useBN
         self.seed = seed
-        paddle.seed(seed)
 
         self.weight = self.create_parameter(
             shape=[in_features, out_features],
@@ -337,9 +335,7 @@ def __init__(
             dtype=precision,
             is_bias=True,
             attr=paddle.ParamAttr(trainable=trainable),
-            default_initializer=paddle.nn.initializer.Normal(
-                mean=bavg if isinstance(bavg, float) else bavg[0], std=stddev
-            ),
+            default_initializer=paddle.nn.initializer.Normal(mean=bavg, std=stddev),
         )
         if self.activation_fn is not None and self.use_timestep:
             self.idt = self.create_parameter(
@@ -414,7 +410,6 @@ def __init__(
         self.activation_fn = activation_fn
         self.resnet_dt = resnet_dt
         self.seed = seed
-        paddle.seed(seed)
 
         outputs_size = self.outputs_size
         weight = []
@@ -434,7 +429,7 @@ def __init__(
             )
             bias.append(
                 self.create_parameter(
-                    shape=[1, outputs_size[ii]],
+                    shape=[outputs_size[ii]],
                     dtype=precision,
                     is_bias=True,
                     attr=paddle.ParamAttr(trainable=trainable),
@@ -477,12 +472,12 @@ def forward(self, xx):
 
             if outputs_size[ii] == outputs_size[ii - 1]:
                 if self.resnet_dt:
-                    xx += hidden * self.idt[ii]
+                    xx += hidden * self.idt[ii - 1]
                 else:
                     xx += hidden
             elif outputs_size[ii] == outputs_size[ii - 1] * 2:
                 if self.resnet_dt:
-                    xx = paddle.concat([xx, xx], axis=1) + hidden * self.idt[ii]
+                    xx = paddle.concat([xx, xx], axis=1) + hidden * self.idt[ii - 1]
                 else:
                     xx = paddle.concat([xx, xx], axis=1) + hidden
             else:
diff --git a/deepmd/utils/spin.py b/deepmd/utils/spin.py
index 38801fe3b9..da99b3e49e 100644
--- a/deepmd/utils/spin.py
+++ b/deepmd/utils/spin.py
@@ -3,12 +3,11 @@
 )
 
 from deepmd.env import (
-    GLOBAL_TF_FLOAT_PRECISION,
-    tf,
+    paddle,
 )
 
 
-class Spin:
+class Spin(paddle.nn.Layer):
     """Class for spin.
 
     Parameters
@@ -27,46 +26,24 @@ def __init__(
         spin_norm: List[float] = None,
         virtual_len: List[float] = None,
     ) -> None:
+        super().__init__()
         """Constructor."""
         self.use_spin = use_spin
         self.spin_norm = spin_norm
         self.virtual_len = virtual_len
         self.ntypes_spin = self.use_spin.count(True)
-
-    def build(
-        self,
-        reuse=None,
-        suffix="",
-    ):
-        """Build the computational graph for the spin.
-
-        Parameters
-        ----------
-        reuse
-            The weights in the networks should be reused when get the variable.
-        suffix
-            Name suffix to identify this descriptor
-
-        Returns
-        -------
-        embedded_types
-            The computational graph for embedded types
-        """
-        name = "spin_attr" + suffix
-        with tf.variable_scope(name, reuse=reuse):
-            t_ntypes_spin = tf.constant(
-                self.ntypes_spin, name="ntypes_spin", dtype=tf.int32
-            )
-            t_virtual_len = tf.constant(
-                self.virtual_len,
-                name="virtual_len",
-                dtype=GLOBAL_TF_FLOAT_PRECISION,
-            )
-            t_spin_norm = tf.constant(
-                self.spin_norm,
-                name="spin_norm",
-                dtype=GLOBAL_TF_FLOAT_PRECISION,
-            )
+        self.register_buffer(
+            "buffer_ntypes_spin",
+            paddle.to_tensor([self.ntypes_spin], dtype="int32"),
+        )
+        self.register_buffer(
+            "buffer_virtual_len",
+            paddle.to_tensor([self.virtual_len], dtype=paddle.get_default_dtype()),
+        )
+        self.register_buffer(
+            "buffer_spin_norm",
+            paddle.to_tensor([self.spin_norm], dtype=paddle.get_default_dtype()),
+        )
 
     def get_ntypes_spin(self) -> int:
         """Returns the number of atom types which contain spin."""
diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc
index c76d339f5b..9d62da78f8 100644
--- a/source/api_cc/src/DeepPot.cc
+++ b/source/api_cc/src/DeepPot.cc
@@ -719,7 +719,7 @@ void DeepPot::init(const std::string& model,
   DPGetDeviceCount(gpu_num);  // check current device environment
 #endif                        // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   if (use_paddle_inference) {
-    /*
+    /* water se_e2_a
     tensorflow::DT_DOUBLE = 2
     tensorflow::DT_FLOAT = 1
     paddle_infer::DataType::FLOAT64 = 7
@@ -751,6 +751,31 @@ void DeepPot::init(const std::string& model,
     static_model is: generated_tensor_9 [fitting.buffer_daparam, []] generated
     name in static_model is: generated_tensor_10
     **/
+    /* spin se_e2_a
+    [buffer_tmap, [4]] generated name in static_model is: generated_tensor_14
+    [buffer_model_type, [4]] generated name in static_model is:
+    generated_tensor_15 [buffer_model_version, [1]] generated name in
+    static_model is: generated_tensor_16 [descrpt.buffer_rcut, []] generated
+    name in static_model is: generated_tensor_3 [descrpt.buffer_ntypes, []]
+    generated name in static_model is: generated_tensor_4 [descrpt.avg_zero, [3,
+    720]] generated name in static_model is: eager_tmp_0 [descrpt.std_ones, [3,
+    720]] generated name in static_model is: eager_tmp_1 [descrpt.t_rcut, []]
+    generated name in static_model is: generated_tensor_5 [descrpt.buffer_sel,
+    [3]] generated name in static_model is: generated_tensor_6
+    [descrpt.buffer_ndescrpt, []] generated name in static_model is:
+    generated_tensor_7 [descrpt.buffer_original_sel, [3]] generated name in
+    static_model is: generated_tensor_8 [descrpt.t_avg, [3, 720]] generated name
+    in static_model is: generated_tensor_9 [descrpt.t_std, [3, 720]] generated
+    name in static_model is: generated_tensor_10
+    [descrpt.spin.buffer_ntypes_spin, [1]] generated name in static_model is:
+    generated_tensor_0 [descrpt.spin.buffer_virtual_len, [1, 1]] generated name
+    in static_model is: generated_tensor_1 [descrpt.spin.buffer_spin_norm, [1,
+    1]] generated name in static_model is: generated_tensor_2
+    [fitting.buffer_dfparam, []] generated name in static_model is:
+    generated_tensor_11 [fitting.buffer_daparam, []] generated name in
+    static_model is: generated_tensor_12 [fitting.t_bias_atom_e, [2]] generated
+    name in static_model is: generated_tensor_13
+   */
     model_version = paddle_get_scalar<std::string>("generated_tensor_14");
     dtype = predictor_get_dtype(predictor, "generated_tensor_0");
     if (dtype == paddle_infer::DataType::FLOAT64) {
@@ -759,12 +784,11 @@ void DeepPot::init(const std::string& model,
       rcut = paddle_get_scalar<float>("generated_tensor_0");
     }
     ntypes = paddle_get_scalar<int32_t>("generated_tensor_2");
-    // ntypes_spin = paddle_get_scalar<int64_t>("buffer_ntypes_spin");
     ntypes_spin = 0;
     dfparam = paddle_get_scalar<int64_t>("generated_tensor_9");
     daparam = paddle_get_scalar<int64_t>("generated_tensor_10");
     model_type = paddle_get_scalar<std::string>("generated_tensor_13");
-    ;
+
     inited = true;
     init_nbor = false;
     return;
diff --git a/source/lib/paddle_src/custom_op_test.py b/source/lib/paddle_src/custom_op_test.py
index 18403c22c5..29ac501935 100644
--- a/source/lib/paddle_src/custom_op_test.py
+++ b/source/lib/paddle_src/custom_op_test.py
@@ -36,10 +36,10 @@ def test_neighbor_stat(place="cpu"):
 
     rcut = 6.0
 
-    coord = paddle.to_tensor(coord, "float32", place=place)
+    coord = paddle.to_tensor(coord, "float64", place=place)
     type = paddle.to_tensor(type, "int32", place=place)
-    natoms = paddle.to_tensor(natoms, "int64", place=place)
-    box = paddle.to_tensor(box, "float32", place=place)
+    natoms = paddle.to_tensor(natoms, "int32", place=place)
+    box = paddle.to_tensor(box, "float64", place=place)
     default_mesh = paddle.to_tensor(default_mesh, "int32", place=place)
 
     mn, dt = paddle_deepmd_lib.neighbor_stat(
@@ -59,8 +59,8 @@ def test_neighbor_stat(place="cpu"):
     # print(dt.shape, dt.min().item(), dt.max().item(), dt.mean().item(), dt.var().item())
     # print(dt_load.shape, dt_load.min().item(), dt_load.max().item(), dt_load.mean().item(), dt_load.var().item())
 
-    print(np.allclose(mn.numpy(), mn_load, 1e-1))
-    print(np.allclose(dt.numpy(), dt_load, 1e-1))
+    print(np.allclose(mn.numpy(), mn_load))
+    print(np.allclose(dt.numpy(), dt_load))
 
 
 def test_prod_env_mat_a(place="cpu"):
@@ -230,11 +230,7 @@ def test_prod_virial_se_a(place="cpu"):
     descrpt_deriv = paddle.to_tensor(descrpt_deriv, place=place)
     rij = paddle.to_tensor(rij, place=place)
     nlist = paddle.to_tensor(nlist, place=place)
-    # natoms = natoms.tolist()
-    natoms = paddle.to_tensor(natoms, place="cpu")  # [192, 192, 64 , 128]
-    # print(natoms.place)
-    # print(natoms.dtype)
-    # print(natoms)
+    natoms = paddle.to_tensor(natoms, place="cpu")
     virial, atom_virial = paddle_deepmd_lib.prod_virial_se_a(
         net_deriv_reshape,
         descrpt_deriv,
diff --git a/source/lib/paddle_src/paddle_neighbor_stat.cc b/source/lib/paddle_src/paddle_neighbor_stat.cc
index 0262a7bdd0..2954c86ae8 100644
--- a/source/lib/paddle_src/paddle_neighbor_stat.cc
+++ b/source/lib/paddle_src/paddle_neighbor_stat.cc
@@ -5,16 +5,49 @@
 
 #undef PADDLE_WITH_CUDA
 #define CHECK_INPUT_CPU(x) PD_CHECK(x.is_cpu(), #x " must be a CPU Tensor.")
+#define CHECK_INPUT_FP64(x)                        \
+  PD_CHECK(x.dtype() == paddle::DataType::FLOAT64, \
+           #x " dtype should be FLOAT64.")
+#define CHECK_INPUT_INT32(x) \
+  PD_CHECK(x.dtype() == paddle::DataType::INT32, #x " dtype should be INT32.")
 #define CHECK_INPUT_DIM(x, value) \
   PD_CHECK(x.shape().size() == value, #x "'s dim should be " #value ".")
+
 typedef double boxtensor_t;
 typedef double compute_t;
 
+void neighbor_stat_cpu_forward_kernel(
+    const int& nloc,
+    const std::vector<int>& d_type,
+    const std::vector<std::vector<int> >& d_nlist_r,
+    const int& ntypes,
+    const std::vector<compute_t>& d_coord3,
+    const int& max_nnei,
+    int* max_nbor_size,
+    double* min_nbor_dist) {
+#pragma omp parallel for
+  for (int ii = 0; ii < nloc; ii++) {
+    if (d_type[ii] < 0) continue;  // virtual atom
+    for (int jj = 0; jj < d_nlist_r[ii].size(); jj++) {
+      int type = d_type[d_nlist_r[ii][jj]];
+      if (type < 0) continue;  // virtual atom
+      max_nbor_size[ii * ntypes + type] += 1;
+      compute_t rij[3] = {
+          d_coord3[d_nlist_r[ii][jj] * 3 + 0] - d_coord3[ii * 3 + 0],
+          d_coord3[d_nlist_r[ii][jj] * 3 + 1] - d_coord3[ii * 3 + 1],
+          d_coord3[d_nlist_r[ii][jj] * 3 + 2] - d_coord3[ii * 3 + 2]};
+      min_nbor_dist[ii * max_nnei + jj] =
+          sqrt(rij[0] * rij[0] + rij[1] * rij[1] + rij[2] * rij[2]);
+    }
+  }
+}
+
+// template<typename FPTYPE>
 std::vector<paddle::Tensor> NeighborStatOpCPUForward(
-    const paddle::Tensor& coord_tensor /*fp32/64*/,
+    const paddle::Tensor& coord_tensor /*fp64*/,
     const paddle::Tensor& type_tensor /*int32*/,
     const paddle::Tensor& natoms_tensor /*int64*/,
-    const paddle::Tensor& box_tensor /*fp32/64*/,
+    const paddle::Tensor& box_tensor /*fp64*/,
     const paddle::Tensor& mesh_tensor /*int32*/,
     const float& rcut) {
   CHECK_INPUT_CPU(coord_tensor);
@@ -23,6 +56,11 @@ std::vector<paddle::Tensor> NeighborStatOpCPUForward(
   CHECK_INPUT_CPU(box_tensor);
   CHECK_INPUT_CPU(mesh_tensor);
 
+  CHECK_INPUT_FP64(coord_tensor);
+  CHECK_INPUT_FP64(box_tensor);
+  CHECK_INPUT_INT32(type_tensor);
+  CHECK_INPUT_INT32(mesh_tensor);
+
   CHECK_INPUT_DIM(coord_tensor, 2);
   CHECK_INPUT_DIM(type_tensor, 2);
   CHECK_INPUT_DIM(natoms_tensor, 1);
@@ -31,9 +69,9 @@ std::vector<paddle::Tensor> NeighborStatOpCPUForward(
   PD_CHECK(natoms_tensor.shape()[0] >= 3,
            "number of atoms should be larger than (or equal to) 3");
 
-  const int64_t* natoms = natoms_tensor.data<int64_t>();
-  int64_t nloc = natoms[0];
-  int64_t nall = natoms[1];
+  const int* natoms = natoms_tensor.data<int>();
+  int nloc = natoms[0];
+  int nall = natoms[1];
   int64_t nsamples = coord_tensor.shape()[0];
   int64_t ntypes = natoms_tensor.shape()[0] - 2;
 
@@ -63,9 +101,9 @@ std::vector<paddle::Tensor> NeighborStatOpCPUForward(
   paddle::Tensor max_nbor_size_tensor = paddle::zeros(
       max_nbor_size_shape, type_tensor.dtype(), type_tensor.place());
 
-  const float* coord = coord_tensor.data<float>();
+  const double* coord = coord_tensor.data<double>();
   const int* type = type_tensor.data<int>();
-  const float* box = box_tensor.data<float>();
+  const double* box = box_tensor.data<double>();
   const int* mesh = mesh_tensor.data<int>();
   int* max_nbor_size = max_nbor_size_tensor.data<int>();
 
@@ -136,31 +174,18 @@ std::vector<paddle::Tensor> NeighborStatOpCPUForward(
   std::vector<int64_t> min_nbor_dist_shape = {nloc * MAX_NNEI};
   paddle::Tensor min_nbor_dist_tensor = paddle::full(
       min_nbor_dist_shape, 10000.0, coord_tensor.dtype(), coord_tensor.place());
-  auto* min_nbor_dist = min_nbor_dist_tensor.data<float>();
+  double* min_nbor_dist = min_nbor_dist_tensor.data<double>();
 
-#pragma omp parallel for
-  for (int ii = 0; ii < nloc; ii++) {
-    if (d_type[ii] < 0) continue;  // virtual atom
-    for (int jj = 0; jj < d_nlist_r[ii].size(); jj++) {
-      int type = d_type[d_nlist_r[ii][jj]];
-      if (type < 0) continue;  // virtual atom
-      max_nbor_size[ii * ntypes + type] += 1;
-      compute_t rij[3] = {
-          d_coord3[d_nlist_r[ii][jj] * 3 + 0] - d_coord3[ii * 3 + 0],
-          d_coord3[d_nlist_r[ii][jj] * 3 + 1] - d_coord3[ii * 3 + 1],
-          d_coord3[d_nlist_r[ii][jj] * 3 + 2] - d_coord3[ii * 3 + 2]};
-      min_nbor_dist[ii * MAX_NNEI + jj] =
-          sqrt(rij[0] * rij[0] + rij[1] * rij[1] + rij[2] * rij[2]);
-    }
-  }
+  neighbor_stat_cpu_forward_kernel(nloc, d_type, d_nlist_r, ntypes, d_coord3,
+                                   MAX_NNEI, max_nbor_size, min_nbor_dist);
   return {max_nbor_size_tensor, min_nbor_dist_tensor};
 }
 
 std::vector<paddle::Tensor> NeighborStatForward(
-    const paddle::Tensor& coord_tensor,  /*float32*/
+    const paddle::Tensor& coord_tensor,  /*float64*/
     const paddle::Tensor& type_tensor,   /*int32*/
     const paddle::Tensor& natoms_tensor, /*int64*/
-    const paddle::Tensor& box_tensor,    /*float32*/
+    const paddle::Tensor& box_tensor,    /*float64*/
     const paddle::Tensor& mesh_tensor,   /*int32*/
     float rcut) {
   if (coord_tensor.is_cpu()) {