diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py
index b8c4902d68..990847c1de 100644
--- a/deepmd/dpmodel/atomic_model/base_atomic_model.py
+++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py
@@ -101,7 +101,11 @@ def forward_common_atomic(
         if self.atom_excl is not None:
             atom_mask = self.atom_excl.build_type_exclude_mask(atype)
             for kk in ret_dict.keys():
-                ret_dict[kk] = ret_dict[kk] * atom_mask[:, :, None]
+                out_shape = ret_dict[kk].shape
+                ret_dict[kk] = (
+                    ret_dict[kk].reshape([out_shape[0], out_shape[1], -1])
+                    * atom_mask[:, :, None]
+                ).reshape(out_shape)
             ret_dict["mask"] = atom_mask
 
         return ret_dict
diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py
index 1fd36bd7e8..22267c895a 100644
--- a/deepmd/dpmodel/infer/deep_eval.py
+++ b/deepmd/dpmodel/infer/deep_eval.py
@@ -123,16 +123,16 @@ def get_dim_aparam(self) -> int:
     @property
     def model_type(self) -> Type["DeepEvalWrapper"]:
         """The the evaluator of the model type."""
-        model_type = self.dp.model_output_type()
-        if model_type == "energy":
+        model_output_type = self.dp.model_output_type()
+        if "energy" in model_output_type:
             return DeepPot
-        elif model_type == "dos":
+        elif "dos" in model_output_type:
             return DeepDOS
-        elif model_type == "dipole":
+        elif "dipole" in model_output_type:
             return DeepDipole
-        elif model_type == "polar":
+        elif "polar" in model_output_type:
             return DeepPolar
-        elif model_type == "wfc":
+        elif "wfc" in model_output_type:
             return DeepWFC
         else:
             raise RuntimeError("Unknown model type")
diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py
index e7cc8d9272..95c448442e 100644
--- a/deepmd/dpmodel/model/base_model.py
+++ b/deepmd/dpmodel/model/base_model.py
@@ -89,7 +89,7 @@ def is_aparam_nall(self) -> bool:
             """
 
         @abstractmethod
-        def model_output_type(self) -> str:
+        def model_output_type(self) -> List[str]:
             """Get the output type for the model."""
 
         @abstractmethod
diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py
index f30f6a4021..d1f671c8de 100644
--- a/deepmd/dpmodel/model/make_model.py
+++ b/deepmd/dpmodel/model/make_model.py
@@ -76,7 +76,7 @@ def model_output_def(self):
             """Get the output def for the model."""
             return ModelOutputDef(self.atomic_output_def())
 
-        def model_output_type(self) -> str:
+        def model_output_type(self) -> List[str]:
             """Get the output type for the model."""
             output_def = self.model_output_def()
             var_defs = output_def.var_defs
@@ -85,12 +85,7 @@ def model_output_type(self) -> str:
                 for kk, vv in var_defs.items()
                 if vv.category == OutputVariableCategory.OUT
             ]
-            if len(vars) == 1:
-                return vars[0]
-            elif len(vars) == 0:
-                raise ValueError("No valid output type found")
-            else:
-                raise ValueError(f"Multiple valid output types found: {vars}")
+            return vars
 
         def call(
             self,
diff --git a/deepmd/dpmodel/output_def.py b/deepmd/dpmodel/output_def.py
index d816ed4e84..ac41513246 100644
--- a/deepmd/dpmodel/output_def.py
+++ b/deepmd/dpmodel/output_def.py
@@ -193,6 +193,11 @@ def __init__(
     ):
         self.name = name
         self.shape = list(shape)
+        # jit doesn't support math.prod(self.shape)
+        self.output_size = 1
+        len_shape = len(self.shape)
+        for i in range(len_shape):
+            self.output_size *= self.shape[i]
         self.atomic = atomic
         self.reduciable = reduciable
         self.r_differentiable = r_differentiable
diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py
index f75052166b..bf6a5b0306 100644
--- a/deepmd/pt/infer/deep_eval.py
+++ b/deepmd/pt/infer/deep_eval.py
@@ -148,18 +148,18 @@ def get_dim_aparam(self) -> int:
     @property
     def model_type(self) -> "DeepEvalWrapper":
         """The the evaluator of the model type."""
-        model_type = self.dp.model["Default"].model_output_type()
-        if model_type == "energy":
+        model_output_type = self.dp.model["Default"].model_output_type()
+        if "energy" in model_output_type:
             return DeepPot
-        elif model_type == "dos":
+        elif "dos" in model_output_type:
             return DeepDOS
-        elif model_type == "dipole":
+        elif "dipole" in model_output_type:
             return DeepDipole
-        elif model_type == "polar":
+        elif "polar" in model_output_type:
             return DeepPolar
-        elif model_type == "global_polar":
+        elif "global_polar" in model_output_type:
             return DeepGlobalPolar
-        elif model_type == "wfc":
+        elif "wfc" in model_output_type:
             return DeepWFC
         else:
             raise RuntimeError("Unknown model type")
diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py
index d3a095ce13..d2f6ab9e52 100644
--- a/deepmd/pt/loss/__init__.py
+++ b/deepmd/pt/loss/__init__.py
@@ -8,9 +8,13 @@
 from .loss import (
     TaskLoss,
 )
+from .tensor import (
+    TensorLoss,
+)
 
 __all__ = [
     "DenoiseLoss",
     "EnergyStdLoss",
+    "TensorLoss",
     "TaskLoss",
 ]
diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py
new file mode 100644
index 0000000000..ee42536557
--- /dev/null
+++ b/deepmd/pt/loss/tensor.py
@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    List,
+)
+
+import torch
+
+from deepmd.pt.loss.loss import (
+    TaskLoss,
+)
+from deepmd.pt.utils import (
+    env,
+)
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
+
+
+class TensorLoss(TaskLoss):
+    def __init__(
+        self,
+        tensor_name: str,
+        tensor_size: int,
+        label_name: str,
+        pref_atomic: float = 0.0,
+        pref: float = 0.0,
+        inference=False,
+        **kwargs,
+    ):
+        r"""Construct a loss for local and global tensors.
+
+        Parameters
+        ----------
+        tensor_name : str
+            The name of the tensor in the model predictions to compute the loss.
+        tensor_size : int
+            The size (dimension) of the tensor.
+        label_name : str
+            The name of the tensor in the labels to compute the loss.
+        pref_atomic : float
+            The prefactor of the weight of atomic loss. It should be larger than or equal to 0.
+        pref : float
+            The prefactor of the weight of global loss. It should be larger than or equal to 0.
+        inference : bool
+            If true, it will output all losses found in output, ignoring the pre-factors.
+        **kwargs
+            Other keyword arguments.
+        """
+        super().__init__()
+        self.tensor_name = tensor_name
+        self.tensor_size = tensor_size
+        self.label_name = label_name
+        self.local_weight = pref_atomic
+        self.global_weight = pref
+        self.inference = inference
+
+        assert (
+            self.local_weight >= 0.0 and self.global_weight >= 0.0
+        ), "Can not assign negative weight to `pref` and `pref_atomic`"
+        self.has_local_weight = self.local_weight > 0.0 or inference
+        self.has_global_weight = self.global_weight > 0.0 or inference
+        assert self.has_local_weight or self.has_global_weight, AssertionError(
+            "Can not assian zero weight both to `pref` and `pref_atomic`"
+        )
+
+    def forward(self, model_pred, label, natoms, learning_rate=0.0, mae=False):
+        """Return loss on local and global tensors.
+
+        Parameters
+        ----------
+        model_pred : dict[str, torch.Tensor]
+            Model predictions.
+        label : dict[str, torch.Tensor]
+            Labels.
+        natoms : int
+            The local atom number.
+
+        Returns
+        -------
+        loss: torch.Tensor
+            Loss for model to minimize.
+        more_loss: dict[str, torch.Tensor]
+            Other losses for display.
+        """
+        del learning_rate, mae
+        loss = torch.tensor(0.0, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)
+        more_loss = {}
+        if (
+            self.has_local_weight
+            and self.tensor_name in model_pred
+            and "atomic_" + self.label_name in label
+        ):
+            local_tensor_pred = model_pred[self.tensor_name].reshape(
+                [-1, natoms, self.tensor_size]
+            )
+            local_tensor_label = label["atomic_" + self.label_name].reshape(
+                [-1, natoms, self.tensor_size]
+            )
+            diff = (local_tensor_pred - local_tensor_label).reshape(
+                [-1, self.tensor_size]
+            )
+            if "mask" in model_pred:
+                diff = diff[model_pred["mask"].reshape([-1]).bool()]
+            l2_local_loss = torch.mean(torch.square(diff))
+            if not self.inference:
+                more_loss[f"l2_local_{self.tensor_name}_loss"] = l2_local_loss.detach()
+            loss += self.local_weight * l2_local_loss
+            rmse_local = l2_local_loss.sqrt()
+            more_loss[f"rmse_local_{self.tensor_name}"] = rmse_local.detach()
+        if (
+            self.has_global_weight
+            and "global_" + self.tensor_name in model_pred
+            and self.label_name in label
+        ):
+            global_tensor_pred = model_pred["global_" + self.tensor_name].reshape(
+                [-1, self.tensor_size]
+            )
+            global_tensor_label = label[self.label_name].reshape([-1, self.tensor_size])
+            diff = global_tensor_pred - global_tensor_label
+            if "mask" in model_pred:
+                atom_num = model_pred["mask"].sum(-1, keepdim=True)
+                l2_global_loss = torch.mean(
+                    torch.sum(torch.square(diff) * atom_num, dim=0) / atom_num.sum()
+                )
+                atom_num = torch.mean(atom_num.float())
+            else:
+                atom_num = natoms
+                l2_global_loss = torch.mean(torch.square(diff))
+            if not self.inference:
+                more_loss[
+                    f"l2_global_{self.tensor_name}_loss"
+                ] = l2_global_loss.detach()
+            loss += self.global_weight * l2_global_loss
+            rmse_global = l2_global_loss.sqrt() / atom_num
+            more_loss[f"rmse_global_{self.tensor_name}"] = rmse_global.detach()
+        return loss, more_loss
+
+    @property
+    def label_requirement(self) -> List[DataRequirementItem]:
+        """Return data label requirements needed for this loss calculation."""
+        label_requirement = []
+        if self.has_local_weight:
+            label_requirement.append(
+                DataRequirementItem(
+                    "atomic_" + self.label_name,
+                    ndof=self.tensor_size,
+                    atomic=True,
+                    must=False,
+                    high_prec=False,
+                )
+            )
+        if self.has_global_weight:
+            label_requirement.append(
+                DataRequirementItem(
+                    self.label_name,
+                    ndof=self.tensor_size,
+                    atomic=False,
+                    must=False,
+                    high_prec=False,
+                )
+            )
+        return label_requirement
diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py
index 8827e3f18b..8180c48c81 100644
--- a/deepmd/pt/model/atomic_model/base_atomic_model.py
+++ b/deepmd/pt/model/atomic_model/base_atomic_model.py
@@ -112,7 +112,11 @@ def forward_common_atomic(
         if self.atom_excl is not None:
             atom_mask = self.atom_excl(atype)
             for kk in ret_dict.keys():
-                ret_dict[kk] = ret_dict[kk] * atom_mask[:, :, None]
+                out_shape = ret_dict[kk].shape
+                ret_dict[kk] = (
+                    ret_dict[kk].reshape([out_shape[0], out_shape[1], -1])
+                    * atom_mask[:, :, None]
+                ).reshape(out_shape)
             ret_dict["mask"] = atom_mask
 
         return ret_dict
diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py
index 87eb391a7e..bd354af8d8 100644
--- a/deepmd/pt/model/model/__init__.py
+++ b/deepmd/pt/model/model/__init__.py
@@ -94,7 +94,7 @@ def get_model(model_params):
     fitting_net["type"] = fitting_net.get("type", "ener")
     fitting_net["ntypes"] = descriptor.get_ntypes()
     fitting_net["mixed_types"] = descriptor.mixed_types()
-    fitting_net["embedding_width"] = descriptor.get_dim_out()
+    fitting_net["embedding_width"] = descriptor.get_dim_emb()
     fitting_net["dim_descrpt"] = descriptor.get_dim_out()
     grad_force = "direct" not in fitting_net["type"]
     if not grad_force:
diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py
index 6629541459..8b6f2c47c1 100644
--- a/deepmd/pt/model/model/dipole_model.py
+++ b/deepmd/pt/model/model/dipole_model.py
@@ -50,6 +50,8 @@ def forward(
                     model_predict["atom_virial"] = model_ret["dipole_derv_c"].squeeze(
                         -3
                     )
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord
diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py
index f2af0fff52..dcf1c36e83 100644
--- a/deepmd/pt/model/model/dp_zbl_model.py
+++ b/deepmd/pt/model/model/dp_zbl_model.py
@@ -63,6 +63,8 @@ def forward(
                 model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-3)
         else:
             model_predict["force"] = model_ret["dforce"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
         return model_predict
 
     @torch.jit.export
diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py
index 1a5706dbbf..cd4f78a2e2 100644
--- a/deepmd/pt/model/model/ener_model.py
+++ b/deepmd/pt/model/model/ener_model.py
@@ -52,6 +52,8 @@ def forward(
                     )
             else:
                 model_predict["force"] = model_ret["dforce"]
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord
diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py
index 60b71400fb..f9daa916a8 100644
--- a/deepmd/pt/model/model/make_model.py
+++ b/deepmd/pt/model/model/make_model.py
@@ -75,7 +75,7 @@ def model_output_def(self):
             return ModelOutputDef(self.atomic_output_def())
 
         @torch.jit.export
-        def model_output_type(self) -> str:
+        def model_output_type(self) -> List[str]:
             """Get the output type for the model."""
             output_def = self.model_output_def()
             var_defs = output_def.var_defs
@@ -86,12 +86,7 @@ def model_output_type(self) -> str:
                 # .value is critical for JIT
                 if vv.category == OutputVariableCategory.OUT.value:
                     vars.append(kk)
-            if len(vars) == 1:
-                return vars[0]
-            elif len(vars) == 0:
-                raise ValueError("No valid output type found")
-            else:
-                raise ValueError(f"Multiple valid output types found: {vars}")
+            return vars
 
         # cannot use the name forward. torch script does not work
         def forward_common(
diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py
index d956a0344c..bf430c6706 100644
--- a/deepmd/pt/model/model/polar_model.py
+++ b/deepmd/pt/model/model/polar_model.py
@@ -42,6 +42,8 @@ def forward(
             model_predict = {}
             model_predict["polar"] = model_ret["polar"]
             model_predict["global_polar"] = model_ret["polar_redu"]
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord
diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py
index 7d2dd221db..21372888d6 100644
--- a/deepmd/pt/model/task/dipole.py
+++ b/deepmd/pt/model/task/dipole.py
@@ -157,7 +157,7 @@ def compute_output_stats(
             The path to the stat file.
 
         """
-        raise NotImplementedError
+        pass
 
     def forward(
         self,
diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py
index 00bf049b97..8bf9cc1c90 100644
--- a/deepmd/pt/model/task/ener.py
+++ b/deepmd/pt/model/task/ener.py
@@ -295,7 +295,7 @@ class EnergyFittingNetDirect(Fitting):
     def __init__(
         self,
         ntypes,
-        embedding_width,
+        dim_descrpt,
         neuron,
         bias_atom_e=None,
         out_dim=1,
@@ -315,7 +315,7 @@ def __init__(
         """
         super().__init__()
         self.ntypes = ntypes
-        self.dim_descrpt = embedding_width
+        self.dim_descrpt = dim_descrpt
         self.use_tebd = use_tebd
         self.out_dim = out_dim
         if bias_atom_e is None:
@@ -329,7 +329,7 @@ def __init__(
         for type_i in range(self.ntypes):
             one = ResidualDeep(
                 type_i,
-                embedding_width,
+                dim_descrpt,
                 neuron,
                 0.0,
                 out_dim=out_dim,
@@ -344,7 +344,7 @@ def __init__(
             for type_i in range(self.ntypes):
                 bias_type = 0.0 if self.use_tebd else bias_atom_e[type_i]
                 one = ResidualDeep(
-                    type_i, embedding_width, neuron, bias_type, resnet_dt=resnet_dt
+                    type_i, dim_descrpt, neuron, bias_type, resnet_dt=resnet_dt
                 )
                 filter_layers.append(one)
         self.filter_layers = torch.nn.ModuleList(filter_layers)
diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py
index 9483d1eb4a..fa4f6d7f37 100644
--- a/deepmd/pt/model/task/polarizability.py
+++ b/deepmd/pt/model/task/polarizability.py
@@ -184,7 +184,7 @@ def compute_output_stats(
             The path to the stat file.
 
         """
-        raise NotImplementedError
+        pass
 
     def forward(
         self,
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index b31822d0ee..77e6b1c709 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -22,6 +22,7 @@
 from deepmd.pt.loss import (
     DenoiseLoss,
     EnergyStdLoss,
+    TensorLoss,
 )
 from deepmd.pt.model.model import (
     get_model,
@@ -192,30 +193,30 @@ def get_dataloader_and_buffer(_data, _params):
                 valid_numb_batch,
             )
 
-        def get_single_model(
-            _model_params,
+        def single_model_stat(
+            _model,
+            _data_stat_nbatch,
             _training_data,
             _validation_data,
             _stat_file_path,
             _data_requirement,
         ):
-            model = get_model(deepcopy(_model_params)).to(DEVICE)
             _training_data.add_data_requirement(_data_requirement)
             if _validation_data is not None:
                 _validation_data.add_data_requirement(_data_requirement)
-            if model.get_dim_fparam() > 0:
+            if _model.get_dim_fparam() > 0:
                 fparam_requirement_items = [
                     DataRequirementItem(
-                        "fparam", model.get_dim_fparam(), atomic=False, must=True
+                        "fparam", _model.get_dim_fparam(), atomic=False, must=True
                     )
                 ]
                 _training_data.add_data_requirement(fparam_requirement_items)
                 if _validation_data is not None:
                     _validation_data.add_data_requirement(fparam_requirement_items)
-            if model.get_dim_aparam() > 0:
+            if _model.get_dim_aparam() > 0:
                 aparam_requirement_items = [
                     DataRequirementItem(
-                        "aparam", model.get_dim_aparam(), atomic=True, must=True
+                        "aparam", _model.get_dim_aparam(), atomic=True, must=True
                     )
                 ]
                 _training_data.add_data_requirement(aparam_requirement_items)
@@ -228,16 +229,21 @@ def get_sample():
                     sampled = make_stat_input(
                         _training_data.systems,
                         _training_data.dataloaders,
-                        _model_params.get("data_stat_nbatch", 10),
+                        _data_stat_nbatch,
                     )
                     return sampled
 
-                model.compute_or_load_stat(
+                _model.compute_or_load_stat(
                     sampled_func=get_sample,
                     stat_file_path=_stat_file_path,
                 )
                 if isinstance(_stat_file_path, DPH5Path):
                     _stat_file_path.root.close()
+
+        def get_single_model(
+            _model_params,
+        ):
+            model = get_model(deepcopy(_model_params)).to(DEVICE)
             return model
 
         def get_lr(lr_params):
@@ -248,7 +254,7 @@ def get_lr(lr_params):
             lr_exp = LearningRateExp(**lr_params)
             return lr_exp
 
-        def get_loss(loss_params, start_lr, _ntypes):
+        def get_loss(loss_params, start_lr, _ntypes, _model):
             loss_type = loss_params.get("type", "ener")
             if loss_type == "ener":
                 loss_params["starter_learning_rate"] = start_lr
@@ -256,6 +262,20 @@ def get_loss(loss_params, start_lr, _ntypes):
             elif loss_type == "denoise":
                 loss_params["ntypes"] = _ntypes
                 return DenoiseLoss(**loss_params)
+            elif loss_type == "tensor":
+                model_output_type = _model.model_output_type()
+                if "mask" in model_output_type:
+                    model_output_type.pop(model_output_type.index("mask"))
+                tensor_name = model_output_type[0]
+                loss_params["tensor_name"] = tensor_name
+                loss_params["tensor_size"] = _model.model_output_def()[
+                    tensor_name
+                ].output_size
+                label_name = tensor_name
+                if label_name == "polar":
+                    label_name = "polarizability"
+                loss_params["label_name"] = label_name
+                return TensorLoss(**loss_params)
             else:
                 raise NotImplementedError
 
@@ -277,12 +297,26 @@ def get_loss(loss_params, start_lr, _ntypes):
         else:
             self.opt_type, self.opt_param = get_opt_param(training_params)
 
+        # Model
+        dp_random.seed(training_params["seed"])
+        if not self.multi_task:
+            self.model = get_single_model(
+                model_params,
+            )
+        else:
+            self.model = {}
+            for model_key in self.model_keys:
+                self.model[model_key] = get_single_model(
+                    model_params["model_dict"][model_key],
+                )
+
         # Loss
         if not self.multi_task:
             self.loss = get_loss(
                 config["loss"],
                 config["learning_rate"]["start_lr"],
                 len(model_params["type_map"]),
+                self.model,
             )
         else:
             self.loss = {}
@@ -293,13 +327,16 @@ def get_loss(loss_params, start_lr, _ntypes):
                 else:
                     lr_param = config["learning_rate"]["start_lr"]
                 ntypes = len(model_params["model_dict"][model_key]["type_map"])
-                self.loss[model_key] = get_loss(loss_param, lr_param, ntypes)
+                self.loss[model_key] = get_loss(
+                    loss_param, lr_param, ntypes, self.model[model_key]
+                )
 
-        # Data + Model
+        # Data
         dp_random.seed(training_params["seed"])
         if not self.multi_task:
-            self.model = get_single_model(
-                model_params,
+            single_model_stat(
+                self.model,
+                model_params.get("data_stat_nbatch", 10),
                 training_data,
                 validation_data,
                 stat_file_path,
@@ -327,11 +364,11 @@ def get_loss(loss_params, start_lr, _ntypes):
                 self.validation_dataloader,
                 self.validation_data,
                 self.valid_numb_batch,
-                self.model,
-            ) = {}, {}, {}, {}, {}, {}
+            ) = {}, {}, {}, {}, {}
             for model_key in self.model_keys:
-                self.model[model_key] = get_single_model(
-                    model_params["model_dict"][model_key],
+                single_model_stat(
+                    self.model[model_key],
+                    model_params["model_dict"][model_key].get("data_stat_nbatch", 10),
                     training_data[model_key],
                     validation_data[model_key],
                     stat_file_path[model_key],
diff --git a/examples/water_tensor/dipole/dipole_input_torch.json b/examples/water_tensor/dipole/dipole_input_torch.json
new file mode 100644
index 0000000000..f6903d3334
--- /dev/null
+++ b/examples/water_tensor/dipole/dipole_input_torch.json
@@ -0,0 +1,84 @@
+{
+  "_comment1": " model parameters",
+  "model": {
+    "type_map": [
+      "O",
+      "H"
+    ],
+    "atom_exclude_types": [
+      1
+    ],
+    "descriptor": {
+      "type": "se_e2_a",
+      "sel": [
+        46,
+        92
+      ],
+      "rcut_smth": 3.80,
+      "rcut": 4.00,
+      "neuron": [
+        25,
+        50,
+        100
+      ],
+      "resnet_dt": false,
+      "axis_neuron": 6,
+      "type_one_side": true,
+      "precision": "float64",
+      "seed": 1,
+      "_comment2": " that's all"
+    },
+    "fitting_net": {
+      "type": "dipole",
+      "neuron": [
+        100,
+        100,
+        100
+      ],
+      "resnet_dt": true,
+      "precision": "float64",
+      "seed": 1,
+      "_comment3": " that's all"
+    },
+    "_comment4": " that's all"
+  },
+  "learning_rate": {
+    "type": "exp",
+    "start_lr": 0.01,
+    "decay_steps": 5000,
+    "_comment5": "that's all"
+  },
+  "loss": {
+    "type": "tensor",
+    "pref": 1.0,
+    "pref_atomic": 1.0,
+    "_comment6": " that's all"
+  },
+  "_comment7": " traing controls",
+  "training": {
+    "training_data": {
+      "systems": [
+        "./training_data_reformat/atomic_system",
+        "./training_data_reformat/global_system"
+      ],
+      "batch_size": "auto",
+      "_comment8": "that's all"
+    },
+    "validation_data": {
+      "systems": [
+        "./validation_data_reformat/atomic_system",
+        "./validation_data_reformat/global_system"
+      ],
+      "batch_size": 1,
+      "numb_btch": 3,
+      "_comment9": "that's all"
+    },
+    "numb_steps": 2000,
+    "seed": 10,
+    "disp_file": "lcurve.out",
+    "disp_freq": 100,
+    "save_freq": 1000,
+    "_comment10": "that's all"
+  },
+  "_comment11": "that's all"
+}
diff --git a/examples/water_tensor/dipole/training_data_reformat/atomic_system/nopbc b/examples/water_tensor/dipole/training_data_reformat/atomic_system/nopbc
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/examples/water_tensor/dipole/training_data_reformat/atomic_system/set.000/atomic_dipole.npy b/examples/water_tensor/dipole/training_data_reformat/atomic_system/set.000/atomic_dipole.npy
new file mode 100644
index 0000000000..2cabc71e21
Binary files /dev/null and b/examples/water_tensor/dipole/training_data_reformat/atomic_system/set.000/atomic_dipole.npy differ
diff --git a/examples/water_tensor/dipole/training_data_reformat/atomic_system/set.000/box.npy b/examples/water_tensor/dipole/training_data_reformat/atomic_system/set.000/box.npy
new file mode 100644
index 0000000000..ed03e9c85b
Binary files /dev/null and b/examples/water_tensor/dipole/training_data_reformat/atomic_system/set.000/box.npy differ
diff --git a/examples/water_tensor/dipole/training_data_reformat/atomic_system/set.000/coord.npy b/examples/water_tensor/dipole/training_data_reformat/atomic_system/set.000/coord.npy
new file mode 100644
index 0000000000..ebee6a5611
Binary files /dev/null and b/examples/water_tensor/dipole/training_data_reformat/atomic_system/set.000/coord.npy differ
diff --git a/examples/water_tensor/dipole/training_data_reformat/atomic_system/type.raw b/examples/water_tensor/dipole/training_data_reformat/atomic_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/examples/water_tensor/dipole/training_data_reformat/atomic_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/examples/water_tensor/dipole/training_data_reformat/atomic_system/type_map.raw b/examples/water_tensor/dipole/training_data_reformat/atomic_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/examples/water_tensor/dipole/training_data_reformat/atomic_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/examples/water_tensor/dipole/training_data_reformat/global_system/nopbc b/examples/water_tensor/dipole/training_data_reformat/global_system/nopbc
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/examples/water_tensor/dipole/training_data_reformat/global_system/set.000/box.npy b/examples/water_tensor/dipole/training_data_reformat/global_system/set.000/box.npy
new file mode 100644
index 0000000000..652530cfe8
Binary files /dev/null and b/examples/water_tensor/dipole/training_data_reformat/global_system/set.000/box.npy differ
diff --git a/examples/water_tensor/dipole/training_data_reformat/global_system/set.000/coord.npy b/examples/water_tensor/dipole/training_data_reformat/global_system/set.000/coord.npy
new file mode 100644
index 0000000000..4f6c37e77a
Binary files /dev/null and b/examples/water_tensor/dipole/training_data_reformat/global_system/set.000/coord.npy differ
diff --git a/examples/water_tensor/dipole/training_data_reformat/global_system/set.000/dipole.npy b/examples/water_tensor/dipole/training_data_reformat/global_system/set.000/dipole.npy
new file mode 100644
index 0000000000..c16efad029
Binary files /dev/null and b/examples/water_tensor/dipole/training_data_reformat/global_system/set.000/dipole.npy differ
diff --git a/examples/water_tensor/dipole/training_data_reformat/global_system/type.raw b/examples/water_tensor/dipole/training_data_reformat/global_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/examples/water_tensor/dipole/training_data_reformat/global_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/examples/water_tensor/dipole/training_data_reformat/global_system/type_map.raw b/examples/water_tensor/dipole/training_data_reformat/global_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/examples/water_tensor/dipole/training_data_reformat/global_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/examples/water_tensor/dipole/validation_data_reformat/atomic_system/nopbc b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/nopbc
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/examples/water_tensor/dipole/validation_data_reformat/atomic_system/set.000/atomic_dipole.npy b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/set.000/atomic_dipole.npy
new file mode 100644
index 0000000000..3a59af8138
Binary files /dev/null and b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/set.000/atomic_dipole.npy differ
diff --git a/examples/water_tensor/dipole/validation_data_reformat/atomic_system/set.000/box.npy b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/set.000/box.npy
new file mode 100644
index 0000000000..6ec95b7173
Binary files /dev/null and b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/set.000/box.npy differ
diff --git a/examples/water_tensor/dipole/validation_data_reformat/atomic_system/set.000/coord.npy b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/set.000/coord.npy
new file mode 100644
index 0000000000..0a9f5f4015
Binary files /dev/null and b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/set.000/coord.npy differ
diff --git a/examples/water_tensor/dipole/validation_data_reformat/atomic_system/type.raw b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/examples/water_tensor/dipole/validation_data_reformat/atomic_system/type_map.raw b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/examples/water_tensor/dipole/validation_data_reformat/atomic_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/examples/water_tensor/dipole/validation_data_reformat/global_system/nopbc b/examples/water_tensor/dipole/validation_data_reformat/global_system/nopbc
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/examples/water_tensor/dipole/validation_data_reformat/global_system/set.000/box.npy b/examples/water_tensor/dipole/validation_data_reformat/global_system/set.000/box.npy
new file mode 100644
index 0000000000..382a14b7b6
Binary files /dev/null and b/examples/water_tensor/dipole/validation_data_reformat/global_system/set.000/box.npy differ
diff --git a/examples/water_tensor/dipole/validation_data_reformat/global_system/set.000/coord.npy b/examples/water_tensor/dipole/validation_data_reformat/global_system/set.000/coord.npy
new file mode 100644
index 0000000000..779b44bad5
Binary files /dev/null and b/examples/water_tensor/dipole/validation_data_reformat/global_system/set.000/coord.npy differ
diff --git a/examples/water_tensor/dipole/validation_data_reformat/global_system/set.000/dipole.npy b/examples/water_tensor/dipole/validation_data_reformat/global_system/set.000/dipole.npy
new file mode 100644
index 0000000000..09c433b931
Binary files /dev/null and b/examples/water_tensor/dipole/validation_data_reformat/global_system/set.000/dipole.npy differ
diff --git a/examples/water_tensor/dipole/validation_data_reformat/global_system/type.raw b/examples/water_tensor/dipole/validation_data_reformat/global_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/examples/water_tensor/dipole/validation_data_reformat/global_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/examples/water_tensor/dipole/validation_data_reformat/global_system/type_map.raw b/examples/water_tensor/dipole/validation_data_reformat/global_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/examples/water_tensor/dipole/validation_data_reformat/global_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/examples/water_tensor/polar/polar_input_torch.json b/examples/water_tensor/polar/polar_input_torch.json
new file mode 100644
index 0000000000..b0329ef609
--- /dev/null
+++ b/examples/water_tensor/polar/polar_input_torch.json
@@ -0,0 +1,90 @@
+{
+  "_comment1": " model parameters",
+  "model": {
+    "type_map": [
+      "O",
+      "H"
+    ],
+    "atom_exclude_types": [
+      1
+    ],
+    "data_stat_nbatch": 10,
+    "descriptor": {
+      "type": "se_e2_a",
+      "sel": [
+        46,
+        92
+      ],
+      "rcut_smth": 5.80,
+      "rcut": 6.00,
+      "neuron": [
+        25,
+        50,
+        100
+      ],
+      "resnet_dt": false,
+      "axis_neuron": 16,
+      "type_one_side": true,
+      "precision": "float64",
+      "seed": 1,
+      "_comment2": " that's all"
+    },
+    "fitting_net": {
+      "type": "polar",
+      "fit_diag": false,
+      "neuron": [
+        100,
+        100,
+        100
+      ],
+      "resnet_dt": true,
+      "precision": "float64",
+      "seed": 1,
+      "_comment3": " that's all"
+    },
+    "_comment4": " that's all"
+  },
+
+  "learning_rate": {
+    "type": "exp",
+    "decay_steps": 5000,
+    "start_lr": 0.01,
+    "stop_lr": 3.51e-7,
+    "_comment5": "that's all"
+  },
+  "loss": {
+    "type": "tensor",
+    "pref_atomic": 1.0,
+    "pref": 1.0,
+    "_comment6": "that's all"
+  },
+
+  "_comment7": " traing controls",
+  "training": {
+    "training_data": {
+      "systems": [
+        "./training_data_reformat/atomic_system",
+        "./training_data_reformat/global_system"
+      ],
+      "batch_size": "auto",
+      "_comment8": "that's all"
+    },
+    "validation_data": {
+      "systems": [
+        "./validation_data_reformat/atomic_system",
+        "./validation_data_reformat/global_system"
+      ],
+      "batch_size": 1,
+      "numb_btch": 3,
+      "_comment9": "that's all"
+    },
+    "numb_steps": 2000,
+    "seed": 10,
+    "disp_file": "lcurve.out",
+    "disp_freq": 100,
+    "save_freq": 1000,
+    "_comment10": "that's all"
+  },
+
+  "_comment11": "that's all"
+}
diff --git a/examples/water_tensor/polar/training_data_reformat/atomic_system/set.000/atomic_polarizability.npy b/examples/water_tensor/polar/training_data_reformat/atomic_system/set.000/atomic_polarizability.npy
new file mode 100644
index 0000000000..2aa2cdd4f2
Binary files /dev/null and b/examples/water_tensor/polar/training_data_reformat/atomic_system/set.000/atomic_polarizability.npy differ
diff --git a/examples/water_tensor/polar/training_data_reformat/atomic_system/set.000/box.npy b/examples/water_tensor/polar/training_data_reformat/atomic_system/set.000/box.npy
new file mode 100644
index 0000000000..a0ce7ef9a7
Binary files /dev/null and b/examples/water_tensor/polar/training_data_reformat/atomic_system/set.000/box.npy differ
diff --git a/examples/water_tensor/polar/training_data_reformat/atomic_system/set.000/coord.npy b/examples/water_tensor/polar/training_data_reformat/atomic_system/set.000/coord.npy
new file mode 100644
index 0000000000..baa2c0a7c3
Binary files /dev/null and b/examples/water_tensor/polar/training_data_reformat/atomic_system/set.000/coord.npy differ
diff --git a/examples/water_tensor/polar/training_data_reformat/atomic_system/type.raw b/examples/water_tensor/polar/training_data_reformat/atomic_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/examples/water_tensor/polar/training_data_reformat/atomic_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/examples/water_tensor/polar/training_data_reformat/atomic_system/type_map.raw b/examples/water_tensor/polar/training_data_reformat/atomic_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/examples/water_tensor/polar/training_data_reformat/atomic_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/examples/water_tensor/polar/training_data_reformat/global_system/set.000/box.npy b/examples/water_tensor/polar/training_data_reformat/global_system/set.000/box.npy
new file mode 100644
index 0000000000..652530cfe8
Binary files /dev/null and b/examples/water_tensor/polar/training_data_reformat/global_system/set.000/box.npy differ
diff --git a/examples/water_tensor/polar/training_data_reformat/global_system/set.000/coord.npy b/examples/water_tensor/polar/training_data_reformat/global_system/set.000/coord.npy
new file mode 100644
index 0000000000..4f6c37e77a
Binary files /dev/null and b/examples/water_tensor/polar/training_data_reformat/global_system/set.000/coord.npy differ
diff --git a/examples/water_tensor/polar/training_data_reformat/global_system/set.000/polarizability.npy b/examples/water_tensor/polar/training_data_reformat/global_system/set.000/polarizability.npy
new file mode 100644
index 0000000000..893767e565
Binary files /dev/null and b/examples/water_tensor/polar/training_data_reformat/global_system/set.000/polarizability.npy differ
diff --git a/examples/water_tensor/polar/training_data_reformat/global_system/type.raw b/examples/water_tensor/polar/training_data_reformat/global_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/examples/water_tensor/polar/training_data_reformat/global_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/examples/water_tensor/polar/training_data_reformat/global_system/type_map.raw b/examples/water_tensor/polar/training_data_reformat/global_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/examples/water_tensor/polar/training_data_reformat/global_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/examples/water_tensor/polar/validation_data_reformat/atomic_system/set.000/atomic_polarizability.npy b/examples/water_tensor/polar/validation_data_reformat/atomic_system/set.000/atomic_polarizability.npy
new file mode 100644
index 0000000000..4cccd9c81a
Binary files /dev/null and b/examples/water_tensor/polar/validation_data_reformat/atomic_system/set.000/atomic_polarizability.npy differ
diff --git a/examples/water_tensor/polar/validation_data_reformat/atomic_system/set.000/box.npy b/examples/water_tensor/polar/validation_data_reformat/atomic_system/set.000/box.npy
new file mode 100644
index 0000000000..a3809a0db5
Binary files /dev/null and b/examples/water_tensor/polar/validation_data_reformat/atomic_system/set.000/box.npy differ
diff --git a/examples/water_tensor/polar/validation_data_reformat/atomic_system/set.000/coord.npy b/examples/water_tensor/polar/validation_data_reformat/atomic_system/set.000/coord.npy
new file mode 100644
index 0000000000..4e18cab336
Binary files /dev/null and b/examples/water_tensor/polar/validation_data_reformat/atomic_system/set.000/coord.npy differ
diff --git a/examples/water_tensor/polar/validation_data_reformat/atomic_system/type.raw b/examples/water_tensor/polar/validation_data_reformat/atomic_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/examples/water_tensor/polar/validation_data_reformat/atomic_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/examples/water_tensor/polar/validation_data_reformat/atomic_system/type_map.raw b/examples/water_tensor/polar/validation_data_reformat/atomic_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/examples/water_tensor/polar/validation_data_reformat/atomic_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/examples/water_tensor/polar/validation_data_reformat/global_system/set.000/box.npy b/examples/water_tensor/polar/validation_data_reformat/global_system/set.000/box.npy
new file mode 100644
index 0000000000..6ec95b7173
Binary files /dev/null and b/examples/water_tensor/polar/validation_data_reformat/global_system/set.000/box.npy differ
diff --git a/examples/water_tensor/polar/validation_data_reformat/global_system/set.000/coord.npy b/examples/water_tensor/polar/validation_data_reformat/global_system/set.000/coord.npy
new file mode 100644
index 0000000000..0a9f5f4015
Binary files /dev/null and b/examples/water_tensor/polar/validation_data_reformat/global_system/set.000/coord.npy differ
diff --git a/examples/water_tensor/polar/validation_data_reformat/global_system/set.000/polarizability.npy b/examples/water_tensor/polar/validation_data_reformat/global_system/set.000/polarizability.npy
new file mode 100644
index 0000000000..cef15d6677
Binary files /dev/null and b/examples/water_tensor/polar/validation_data_reformat/global_system/set.000/polarizability.npy differ
diff --git a/examples/water_tensor/polar/validation_data_reformat/global_system/type.raw b/examples/water_tensor/polar/validation_data_reformat/global_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/examples/water_tensor/polar/validation_data_reformat/global_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/examples/water_tensor/polar/validation_data_reformat/global_system/type_map.raw b/examples/water_tensor/polar/validation_data_reformat/global_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/examples/water_tensor/polar/validation_data_reformat/global_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py
index 49abcf2f90..647bee2bbb 100644
--- a/source/tests/common/test_examples.py
+++ b/source/tests/common/test_examples.py
@@ -33,6 +33,8 @@
     p_examples / "nopbc" / "train" / "input.json",
     p_examples / "water_tensor" / "dipole" / "dipole_input.json",
     p_examples / "water_tensor" / "polar" / "polar_input.json",
+    p_examples / "water_tensor" / "dipole" / "dipole_input_torch.json",
+    p_examples / "water_tensor" / "polar" / "polar_input_torch.json",
     p_examples / "water_multi_task" / "ener_dipole" / "input.json",
     p_examples / "fparam" / "train" / "input.json",
     p_examples / "fparam" / "train" / "input_aparam.json",
diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py
index 13e47a953b..0a052103b6 100644
--- a/source/tests/pt/test_training.py
+++ b/source/tests/pt/test_training.py
@@ -34,10 +34,25 @@ def test_trainable(self):
         fix_params = deepcopy(self.config)
         fix_params["model"]["descriptor"]["trainable"] = False
         fix_params["model"]["fitting_net"]["trainable"] = False
-        trainer_fix = get_trainer(fix_params)
-        model_dict_before_training = deepcopy(trainer_fix.model.state_dict())
-        trainer_fix.run()
-        model_dict_after_training = deepcopy(trainer_fix.model.state_dict())
+        free_descriptor = hasattr(self, "not_all_grad") and self.not_all_grad
+        if free_descriptor:
+            # can not set requires_grad false for all parameters,
+            # because the input coord has no grad, thus the loss if all set to false
+            # we only check trainable for fitting net
+            fix_params["model"]["descriptor"]["trainable"] = True
+            trainer_fix = get_trainer(fix_params)
+            model_dict_before_training = deepcopy(
+                trainer_fix.model.fitting_net.state_dict()
+            )
+            trainer_fix.run()
+            model_dict_after_training = deepcopy(
+                trainer_fix.model.fitting_net.state_dict()
+            )
+        else:
+            trainer_fix = get_trainer(fix_params)
+            model_dict_before_training = deepcopy(trainer_fix.model.state_dict())
+            trainer_fix.run()
+            model_dict_after_training = deepcopy(trainer_fix.model.state_dict())
         for key in model_dict_before_training:
             torch.testing.assert_close(
                 model_dict_before_training[key], model_dict_after_training[key]
@@ -141,5 +156,191 @@ def tearDown(self) -> None:
         DPTrainTest.tearDown(self)
 
 
+class TestDipoleModelSeA(unittest.TestCase, DPTrainTest):
+    def setUp(self):
+        input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        data_file_atomic = str(
+            Path(__file__).parent / "water_tensor/dipole/atomic_system"
+        )
+        data_file_global = str(
+            Path(__file__).parent / "water_tensor/dipole/global_system"
+        )
+        self.config["training"]["training_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["training"]["validation_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["model"] = deepcopy(model_se_e2_a)
+        self.config["model"]["atom_exclude_types"] = [1]
+        self.config["model"]["fitting_net"]["type"] = "dipole"
+        self.config["training"]["numb_steps"] = 1
+        self.config["training"]["save_freq"] = 1
+
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
+
+class TestDipoleModelDPA1(unittest.TestCase, DPTrainTest):
+    def setUp(self):
+        input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        data_file_atomic = str(
+            Path(__file__).parent / "water_tensor/dipole/atomic_system"
+        )
+        data_file_global = str(
+            Path(__file__).parent / "water_tensor/dipole/global_system"
+        )
+        self.config["training"]["training_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["training"]["validation_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["model"] = deepcopy(model_dpa1)
+        self.config["model"]["atom_exclude_types"] = [1]
+        self.config["model"]["fitting_net"]["type"] = "dipole"
+        self.config["training"]["numb_steps"] = 1
+        self.config["training"]["save_freq"] = 1
+
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
+
+class TestDipoleModelDPA2(unittest.TestCase, DPTrainTest):
+    def setUp(self):
+        input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        data_file_atomic = str(
+            Path(__file__).parent / "water_tensor/dipole/atomic_system"
+        )
+        data_file_global = str(
+            Path(__file__).parent / "water_tensor/dipole/global_system"
+        )
+        self.config["training"]["training_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["training"]["validation_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["model"] = deepcopy(model_dpa2)
+        self.config["model"]["atom_exclude_types"] = [1]
+        self.config["model"]["fitting_net"]["type"] = "dipole"
+        self.config["training"]["numb_steps"] = 1
+        self.config["training"]["save_freq"] = 1
+
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
+
+class TestPolarModelSeA(unittest.TestCase, DPTrainTest):
+    def setUp(self):
+        input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        data_file_atomic = str(
+            Path(__file__).parent / "water_tensor/polar/atomic_system"
+        )
+        data_file_global = str(
+            Path(__file__).parent / "water_tensor/polar/global_system"
+        )
+        self.config["training"]["training_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["training"]["validation_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["model"] = deepcopy(model_se_e2_a)
+        self.config["model"]["atom_exclude_types"] = [1]
+        self.config["model"]["fitting_net"]["type"] = "polar"
+        self.config["model"]["fitting_net"]["fit_diag"] = False
+        self.config["training"]["numb_steps"] = 1
+        self.config["training"]["save_freq"] = 1
+        # can not set requires_grad false for all parameters,
+        # because the input coord has no grad, thus the loss if all set to false
+        self.not_all_grad = True
+
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
+
+class TestPolarModelDPA1(unittest.TestCase, DPTrainTest):
+    def setUp(self):
+        input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        data_file_atomic = str(
+            Path(__file__).parent / "water_tensor/polar/atomic_system"
+        )
+        data_file_global = str(
+            Path(__file__).parent / "water_tensor/polar/global_system"
+        )
+        self.config["training"]["training_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["training"]["validation_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["model"] = deepcopy(model_dpa1)
+        self.config["model"]["atom_exclude_types"] = [1]
+        self.config["model"]["fitting_net"]["type"] = "polar"
+        self.config["model"]["fitting_net"]["fit_diag"] = False
+        self.config["training"]["numb_steps"] = 1
+        self.config["training"]["save_freq"] = 1
+        # can not set requires_grad false for all parameters,
+        # because the input coord has no grad, thus the loss if all set to false
+        self.not_all_grad = True
+
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
+
+class TestPolarModelDPA2(unittest.TestCase, DPTrainTest):
+    def setUp(self):
+        input_json = str(Path(__file__).parent / "water_tensor/se_e2_a.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        data_file_atomic = str(
+            Path(__file__).parent / "water_tensor/polar/atomic_system"
+        )
+        data_file_global = str(
+            Path(__file__).parent / "water_tensor/polar/global_system"
+        )
+        self.config["training"]["training_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["training"]["validation_data"]["systems"] = [
+            data_file_atomic,
+            data_file_global,
+        ]
+        self.config["model"] = deepcopy(model_dpa2)
+        self.config["model"]["atom_exclude_types"] = [1]
+        self.config["model"]["fitting_net"]["type"] = "polar"
+        self.config["model"]["fitting_net"]["fit_diag"] = False
+        self.config["training"]["numb_steps"] = 1
+        self.config["training"]["save_freq"] = 1
+        # can not set requires_grad false for all parameters,
+        # because the input coord has no grad, thus the loss if all set to false
+        self.not_all_grad = True
+
+    def tearDown(self) -> None:
+        DPTrainTest.tearDown(self)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt/water_tensor/dipole/atomic_system/nopbc b/source/tests/pt/water_tensor/dipole/atomic_system/nopbc
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/source/tests/pt/water_tensor/dipole/atomic_system/set.000/atomic_dipole.npy b/source/tests/pt/water_tensor/dipole/atomic_system/set.000/atomic_dipole.npy
new file mode 100644
index 0000000000..2cabc71e21
Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/atomic_system/set.000/atomic_dipole.npy differ
diff --git a/source/tests/pt/water_tensor/dipole/atomic_system/set.000/box.npy b/source/tests/pt/water_tensor/dipole/atomic_system/set.000/box.npy
new file mode 100644
index 0000000000..ed03e9c85b
Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/atomic_system/set.000/box.npy differ
diff --git a/source/tests/pt/water_tensor/dipole/atomic_system/set.000/coord.npy b/source/tests/pt/water_tensor/dipole/atomic_system/set.000/coord.npy
new file mode 100644
index 0000000000..ebee6a5611
Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/atomic_system/set.000/coord.npy differ
diff --git a/source/tests/pt/water_tensor/dipole/atomic_system/type.raw b/source/tests/pt/water_tensor/dipole/atomic_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/source/tests/pt/water_tensor/dipole/atomic_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/source/tests/pt/water_tensor/dipole/atomic_system/type_map.raw b/source/tests/pt/water_tensor/dipole/atomic_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/source/tests/pt/water_tensor/dipole/atomic_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/source/tests/pt/water_tensor/dipole/global_system/nopbc b/source/tests/pt/water_tensor/dipole/global_system/nopbc
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/source/tests/pt/water_tensor/dipole/global_system/set.000/box.npy b/source/tests/pt/water_tensor/dipole/global_system/set.000/box.npy
new file mode 100644
index 0000000000..652530cfe8
Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/global_system/set.000/box.npy differ
diff --git a/source/tests/pt/water_tensor/dipole/global_system/set.000/coord.npy b/source/tests/pt/water_tensor/dipole/global_system/set.000/coord.npy
new file mode 100644
index 0000000000..4f6c37e77a
Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/global_system/set.000/coord.npy differ
diff --git a/source/tests/pt/water_tensor/dipole/global_system/set.000/dipole.npy b/source/tests/pt/water_tensor/dipole/global_system/set.000/dipole.npy
new file mode 100644
index 0000000000..c16efad029
Binary files /dev/null and b/source/tests/pt/water_tensor/dipole/global_system/set.000/dipole.npy differ
diff --git a/source/tests/pt/water_tensor/dipole/global_system/type.raw b/source/tests/pt/water_tensor/dipole/global_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/source/tests/pt/water_tensor/dipole/global_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/source/tests/pt/water_tensor/dipole/global_system/type_map.raw b/source/tests/pt/water_tensor/dipole/global_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/source/tests/pt/water_tensor/dipole/global_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/source/tests/pt/water_tensor/polar/atomic_system/set.000/atomic_polarizability.npy b/source/tests/pt/water_tensor/polar/atomic_system/set.000/atomic_polarizability.npy
new file mode 100644
index 0000000000..2aa2cdd4f2
Binary files /dev/null and b/source/tests/pt/water_tensor/polar/atomic_system/set.000/atomic_polarizability.npy differ
diff --git a/source/tests/pt/water_tensor/polar/atomic_system/set.000/box.npy b/source/tests/pt/water_tensor/polar/atomic_system/set.000/box.npy
new file mode 100644
index 0000000000..a0ce7ef9a7
Binary files /dev/null and b/source/tests/pt/water_tensor/polar/atomic_system/set.000/box.npy differ
diff --git a/source/tests/pt/water_tensor/polar/atomic_system/set.000/coord.npy b/source/tests/pt/water_tensor/polar/atomic_system/set.000/coord.npy
new file mode 100644
index 0000000000..baa2c0a7c3
Binary files /dev/null and b/source/tests/pt/water_tensor/polar/atomic_system/set.000/coord.npy differ
diff --git a/source/tests/pt/water_tensor/polar/atomic_system/type.raw b/source/tests/pt/water_tensor/polar/atomic_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/source/tests/pt/water_tensor/polar/atomic_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/source/tests/pt/water_tensor/polar/atomic_system/type_map.raw b/source/tests/pt/water_tensor/polar/atomic_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/source/tests/pt/water_tensor/polar/atomic_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/source/tests/pt/water_tensor/polar/global_system/set.000/box.npy b/source/tests/pt/water_tensor/polar/global_system/set.000/box.npy
new file mode 100644
index 0000000000..652530cfe8
Binary files /dev/null and b/source/tests/pt/water_tensor/polar/global_system/set.000/box.npy differ
diff --git a/source/tests/pt/water_tensor/polar/global_system/set.000/coord.npy b/source/tests/pt/water_tensor/polar/global_system/set.000/coord.npy
new file mode 100644
index 0000000000..4f6c37e77a
Binary files /dev/null and b/source/tests/pt/water_tensor/polar/global_system/set.000/coord.npy differ
diff --git a/source/tests/pt/water_tensor/polar/global_system/set.000/polarizability.npy b/source/tests/pt/water_tensor/polar/global_system/set.000/polarizability.npy
new file mode 100644
index 0000000000..893767e565
Binary files /dev/null and b/source/tests/pt/water_tensor/polar/global_system/set.000/polarizability.npy differ
diff --git a/source/tests/pt/water_tensor/polar/global_system/type.raw b/source/tests/pt/water_tensor/polar/global_system/type.raw
new file mode 100644
index 0000000000..6c71c85e58
--- /dev/null
+++ b/source/tests/pt/water_tensor/polar/global_system/type.raw
@@ -0,0 +1 @@
+0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1 0 1 1
diff --git a/source/tests/pt/water_tensor/polar/global_system/type_map.raw b/source/tests/pt/water_tensor/polar/global_system/type_map.raw
new file mode 100644
index 0000000000..e900768b1d
--- /dev/null
+++ b/source/tests/pt/water_tensor/polar/global_system/type_map.raw
@@ -0,0 +1,2 @@
+O
+H
diff --git a/source/tests/pt/water_tensor/se_e2_a.json b/source/tests/pt/water_tensor/se_e2_a.json
new file mode 100644
index 0000000000..e53caafc96
--- /dev/null
+++ b/source/tests/pt/water_tensor/se_e2_a.json
@@ -0,0 +1,85 @@
+{
+  "_comment1": " model parameters",
+  "model": {
+    "type_map": [
+      "O",
+      "H"
+    ],
+    "atom_exclude_types": [
+      1
+    ],
+    "descriptor": {
+      "type": "se_e2_a",
+      "sel": [
+        46,
+        92
+      ],
+      "rcut_smth": 3.80,
+      "rcut": 4.00,
+      "neuron": [
+        25,
+        50,
+        100
+      ],
+      "resnet_dt": false,
+      "axis_neuron": 6,
+      "type_one_side": true,
+      "precision": "float64",
+      "seed": 1,
+      "_comment2": " that's all"
+    },
+    "fitting_net": {
+      "type": "dipole",
+      "neuron": [
+        100,
+        100,
+        100
+      ],
+      "resnet_dt": true,
+      "precision": "float64",
+      "seed": 1,
+      "_comment3": " that's all"
+    },
+    "_comment4": " that's all"
+  },
+  "learning_rate": {
+    "type": "exp",
+    "start_lr": 0.001,
+    "stop_lr": 3.51e-8,
+    "decay_steps": 5000,
+    "_comment5": "that's all"
+  },
+  "loss": {
+    "type": "tensor",
+    "pref": 1.0,
+    "pref_atomic": 1.0,
+    "_comment6": " that's all"
+  },
+  "_comment7": " traing controls",
+  "training": {
+    "training_data": {
+      "systems": [
+        "pt/water_tensor/dipole/atomic_system",
+        "pt/water_tensor/dipole/global_system"
+      ],
+      "batch_size": "auto",
+      "_comment8": "that's all"
+    },
+    "validation_data": {
+      "systems": [
+        "pt/water_tensor/dipole/atomic_system",
+        "pt/water_tensor/dipole/global_system"
+      ],
+      "batch_size": 1,
+      "numb_btch": 3,
+      "_comment9": "that's all"
+    },
+    "numb_steps": 2000,
+    "seed": 10,
+    "disp_file": "lcurve.out",
+    "disp_freq": 100,
+    "save_freq": 1000,
+    "_comment10": "that's all"
+  },
+  "_comment11": "that's all"
+}