diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h
index 3da5c34184..884f76ab6f 100644
--- a/source/api_cc/include/DeepPot.h
+++ b/source/api_cc/include/DeepPot.h
@@ -61,30 +61,31 @@ class DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Request atomic energy and virial if atomic is true.
    * @{
    **/
-  virtual void computew(
-      std::vector<double>& ener,
-      std::vector<double>& force,
-      std::vector<double>& virial,
-      std::vector<double>& atom_energy,
-      std::vector<double>& atom_virial,
-      const std::vector<double>& coord,
-      const std::vector<int>& atype,
-      const std::vector<double>& box,
-      const std::vector<double>& fparam = std::vector<double>(),
-      const std::vector<double>& aparam = std::vector<double>()) = 0;
-  virtual void computew(
-      std::vector<double>& ener,
-      std::vector<float>& force,
-      std::vector<float>& virial,
-      std::vector<float>& atom_energy,
-      std::vector<float>& atom_virial,
-      const std::vector<float>& coord,
-      const std::vector<int>& atype,
-      const std::vector<float>& box,
-      const std::vector<float>& fparam = std::vector<float>(),
-      const std::vector<float>& aparam = std::vector<float>()) = 0;
+  virtual void computew(std::vector<double>& ener,
+                        std::vector<double>& force,
+                        std::vector<double>& virial,
+                        std::vector<double>& atom_energy,
+                        std::vector<double>& atom_virial,
+                        const std::vector<double>& coord,
+                        const std::vector<int>& atype,
+                        const std::vector<double>& box,
+                        const std::vector<double>& fparam,
+                        const std::vector<double>& aparam,
+                        const bool atomic) = 0;
+  virtual void computew(std::vector<double>& ener,
+                        std::vector<float>& force,
+                        std::vector<float>& virial,
+                        std::vector<float>& atom_energy,
+                        std::vector<float>& atom_virial,
+                        const std::vector<float>& coord,
+                        const std::vector<int>& atype,
+                        const std::vector<float>& box,
+                        const std::vector<float>& fparam,
+                        const std::vector<float>& aparam,
+                        const bool atomic) = 0;
   /** @} */
   /**
    * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial
@@ -111,36 +112,37 @@ class DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Request atomic energy and virial if atomic is true.
    * @{
    **/
-  virtual void computew(
-      std::vector<double>& ener,
-      std::vector<double>& force,
-      std::vector<double>& virial,
-      std::vector<double>& atom_energy,
-      std::vector<double>& atom_virial,
-      const std::vector<double>& coord,
-      const std::vector<int>& atype,
-      const std::vector<double>& box,
-      const int nghost,
-      const InputNlist& inlist,
-      const int& ago,
-      const std::vector<double>& fparam = std::vector<double>(),
-      const std::vector<double>& aparam = std::vector<double>()) = 0;
-  virtual void computew(
-      std::vector<double>& ener,
-      std::vector<float>& force,
-      std::vector<float>& virial,
-      std::vector<float>& atom_energy,
-      std::vector<float>& atom_virial,
-      const std::vector<float>& coord,
-      const std::vector<int>& atype,
-      const std::vector<float>& box,
-      const int nghost,
-      const InputNlist& inlist,
-      const int& ago,
-      const std::vector<float>& fparam = std::vector<float>(),
-      const std::vector<float>& aparam = std::vector<float>()) = 0;
+  virtual void computew(std::vector<double>& ener,
+                        std::vector<double>& force,
+                        std::vector<double>& virial,
+                        std::vector<double>& atom_energy,
+                        std::vector<double>& atom_virial,
+                        const std::vector<double>& coord,
+                        const std::vector<int>& atype,
+                        const std::vector<double>& box,
+                        const int nghost,
+                        const InputNlist& inlist,
+                        const int& ago,
+                        const std::vector<double>& fparam,
+                        const std::vector<double>& aparam,
+                        const bool atomic) = 0;
+  virtual void computew(std::vector<double>& ener,
+                        std::vector<float>& force,
+                        std::vector<float>& virial,
+                        std::vector<float>& atom_energy,
+                        std::vector<float>& atom_virial,
+                        const std::vector<float>& coord,
+                        const std::vector<int>& atype,
+                        const std::vector<float>& box,
+                        const int nghost,
+                        const InputNlist& inlist,
+                        const int& ago,
+                        const std::vector<float>& fparam,
+                        const std::vector<float>& aparam,
+                        const bool atomic) = 0;
   /** @} */
 
   /**
@@ -167,32 +169,33 @@ class DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Request atomic energy and virial if atomic is true.
    * @{
    **/
-  virtual void computew_mixed_type(
-      std::vector<double>& ener,
-      std::vector<double>& force,
-      std::vector<double>& virial,
-      std::vector<double>& atom_energy,
-      std::vector<double>& atom_virial,
-      const int& nframes,
-      const std::vector<double>& coord,
-      const std::vector<int>& atype,
-      const std::vector<double>& box,
-      const std::vector<double>& fparam = std::vector<double>(),
-      const std::vector<double>& aparam = std::vector<double>()) = 0;
-  virtual void computew_mixed_type(
-      std::vector<double>& ener,
-      std::vector<float>& force,
-      std::vector<float>& virial,
-      std::vector<float>& atom_energy,
-      std::vector<float>& atom_virial,
-      const int& nframes,
-      const std::vector<float>& coord,
-      const std::vector<int>& atype,
-      const std::vector<float>& box,
-      const std::vector<float>& fparam = std::vector<float>(),
-      const std::vector<float>& aparam = std::vector<float>()) = 0;
+  virtual void computew_mixed_type(std::vector<double>& ener,
+                                   std::vector<double>& force,
+                                   std::vector<double>& virial,
+                                   std::vector<double>& atom_energy,
+                                   std::vector<double>& atom_virial,
+                                   const int& nframes,
+                                   const std::vector<double>& coord,
+                                   const std::vector<int>& atype,
+                                   const std::vector<double>& box,
+                                   const std::vector<double>& fparam,
+                                   const std::vector<double>& aparam,
+                                   const bool atomic) = 0;
+  virtual void computew_mixed_type(std::vector<double>& ener,
+                                   std::vector<float>& force,
+                                   std::vector<float>& virial,
+                                   std::vector<float>& atom_energy,
+                                   std::vector<float>& atom_virial,
+                                   const int& nframes,
+                                   const std::vector<float>& coord,
+                                   const std::vector<int>& atype,
+                                   const std::vector<float>& box,
+                                   const std::vector<float>& fparam,
+                                   const std::vector<float>& aparam,
+                                   const bool atomic) = 0;
   /** @} */
   /**
    * @brief Get the cutoff radius.
diff --git a/source/api_cc/include/DeepPotPT.h b/source/api_cc/include/DeepPotPT.h
index 25072958b9..973c02c434 100644
--- a/source/api_cc/include/DeepPotPT.h
+++ b/source/api_cc/include/DeepPotPT.h
@@ -60,6 +60,7 @@ class DeepPotPT : public DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Whether to compute the atomic energy and virial.
    **/
   template <typename VALUETYPE, typename ENERGYVTYPE>
   void compute(ENERGYVTYPE& ener,
@@ -70,8 +71,9 @@ class DeepPotPT : public DeepPotBase {
                const std::vector<VALUETYPE>& coord,
                const std::vector<int>& atype,
                const std::vector<VALUETYPE>& box,
-               const std::vector<VALUETYPE>& fparam = std::vector<VALUETYPE>(),
-               const std::vector<VALUETYPE>& aparam = std::vector<VALUETYPE>());
+               const std::vector<VALUETYPE>& fparam,
+               const std::vector<VALUETYPE>& aparam,
+               const bool atomic);
   /**
    * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial
    *by using this DP.
@@ -96,6 +98,7 @@ class DeepPotPT : public DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Whether to compute the atomic energy and virial.
    **/
   template <typename VALUETYPE, typename ENERGYVTYPE>
   void compute(ENERGYVTYPE& ener,
@@ -109,8 +112,9 @@ class DeepPotPT : public DeepPotBase {
                const int nghost,
                const InputNlist& lmp_list,
                const int& ago,
-               const std::vector<VALUETYPE>& fparam = std::vector<VALUETYPE>(),
-               const std::vector<VALUETYPE>& aparam = std::vector<VALUETYPE>());
+               const std::vector<VALUETYPE>& fparam,
+               const std::vector<VALUETYPE>& aparam,
+               const bool atomic);
   /**
    * @brief Evaluate the energy, force, and virial with the mixed type
    *by using this DP.
@@ -132,18 +136,19 @@ class DeepPotPT : public DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Whether to compute the atomic energy and virial.
    **/
   template <typename VALUETYPE, typename ENERGYVTYPE>
-  void compute_mixed_type(
-      ENERGYVTYPE& ener,
-      std::vector<VALUETYPE>& force,
-      std::vector<VALUETYPE>& virial,
-      const int& nframes,
-      const std::vector<VALUETYPE>& coord,
-      const std::vector<int>& atype,
-      const std::vector<VALUETYPE>& box,
-      const std::vector<VALUETYPE>& fparam = std::vector<VALUETYPE>(),
-      const std::vector<VALUETYPE>& aparam = std::vector<VALUETYPE>());
+  void compute_mixed_type(ENERGYVTYPE& ener,
+                          std::vector<VALUETYPE>& force,
+                          std::vector<VALUETYPE>& virial,
+                          const int& nframes,
+                          const std::vector<VALUETYPE>& coord,
+                          const std::vector<int>& atype,
+                          const std::vector<VALUETYPE>& box,
+                          const std::vector<VALUETYPE>& fparam,
+                          const std::vector<VALUETYPE>& aparam,
+                          const bool atomic);
   /**
    * @brief Evaluate the energy, force, and virial with the mixed type
    *by using this DP.
@@ -167,20 +172,21 @@ class DeepPotPT : public DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Whether to compute the atomic energy and virial.
    **/
   template <typename VALUETYPE, typename ENERGYVTYPE>
-  void compute_mixed_type(
-      ENERGYVTYPE& ener,
-      std::vector<VALUETYPE>& force,
-      std::vector<VALUETYPE>& virial,
-      std::vector<VALUETYPE>& atom_energy,
-      std::vector<VALUETYPE>& atom_virial,
-      const int& nframes,
-      const std::vector<VALUETYPE>& coord,
-      const std::vector<int>& atype,
-      const std::vector<VALUETYPE>& box,
-      const std::vector<VALUETYPE>& fparam = std::vector<VALUETYPE>(),
-      const std::vector<VALUETYPE>& aparam = std::vector<VALUETYPE>());
+  void compute_mixed_type(ENERGYVTYPE& ener,
+                          std::vector<VALUETYPE>& force,
+                          std::vector<VALUETYPE>& virial,
+                          std::vector<VALUETYPE>& atom_energy,
+                          std::vector<VALUETYPE>& atom_virial,
+                          const int& nframes,
+                          const std::vector<VALUETYPE>& coord,
+                          const std::vector<int>& atype,
+                          const std::vector<VALUETYPE>& box,
+                          const std::vector<VALUETYPE>& fparam,
+                          const std::vector<VALUETYPE>& aparam,
+                          const bool atomic);
 
  public:
   /**
@@ -248,8 +254,9 @@ class DeepPotPT : public DeepPotBase {
                 const std::vector<double>& coord,
                 const std::vector<int>& atype,
                 const std::vector<double>& box,
-                const std::vector<double>& fparam = std::vector<double>(),
-                const std::vector<double>& aparam = std::vector<double>());
+                const std::vector<double>& fparam,
+                const std::vector<double>& aparam,
+                const bool atomic);
   void computew(std::vector<double>& ener,
                 std::vector<float>& force,
                 std::vector<float>& virial,
@@ -258,8 +265,9 @@ class DeepPotPT : public DeepPotBase {
                 const std::vector<float>& coord,
                 const std::vector<int>& atype,
                 const std::vector<float>& box,
-                const std::vector<float>& fparam = std::vector<float>(),
-                const std::vector<float>& aparam = std::vector<float>());
+                const std::vector<float>& fparam,
+                const std::vector<float>& aparam,
+                const bool atomic);
   void computew(std::vector<double>& ener,
                 std::vector<double>& force,
                 std::vector<double>& virial,
@@ -271,8 +279,9 @@ class DeepPotPT : public DeepPotBase {
                 const int nghost,
                 const InputNlist& inlist,
                 const int& ago,
-                const std::vector<double>& fparam = std::vector<double>(),
-                const std::vector<double>& aparam = std::vector<double>());
+                const std::vector<double>& fparam,
+                const std::vector<double>& aparam,
+                const bool atomic);
   void computew(std::vector<double>& ener,
                 std::vector<float>& force,
                 std::vector<float>& virial,
@@ -284,32 +293,33 @@ class DeepPotPT : public DeepPotBase {
                 const int nghost,
                 const InputNlist& inlist,
                 const int& ago,
-                const std::vector<float>& fparam = std::vector<float>(),
-                const std::vector<float>& aparam = std::vector<float>());
-  void computew_mixed_type(
-      std::vector<double>& ener,
-      std::vector<double>& force,
-      std::vector<double>& virial,
-      std::vector<double>& atom_energy,
-      std::vector<double>& atom_virial,
-      const int& nframes,
-      const std::vector<double>& coord,
-      const std::vector<int>& atype,
-      const std::vector<double>& box,
-      const std::vector<double>& fparam = std::vector<double>(),
-      const std::vector<double>& aparam = std::vector<double>());
-  void computew_mixed_type(
-      std::vector<double>& ener,
-      std::vector<float>& force,
-      std::vector<float>& virial,
-      std::vector<float>& atom_energy,
-      std::vector<float>& atom_virial,
-      const int& nframes,
-      const std::vector<float>& coord,
-      const std::vector<int>& atype,
-      const std::vector<float>& box,
-      const std::vector<float>& fparam = std::vector<float>(),
-      const std::vector<float>& aparam = std::vector<float>());
+                const std::vector<float>& fparam,
+                const std::vector<float>& aparam,
+                const bool atomic);
+  void computew_mixed_type(std::vector<double>& ener,
+                           std::vector<double>& force,
+                           std::vector<double>& virial,
+                           std::vector<double>& atom_energy,
+                           std::vector<double>& atom_virial,
+                           const int& nframes,
+                           const std::vector<double>& coord,
+                           const std::vector<int>& atype,
+                           const std::vector<double>& box,
+                           const std::vector<double>& fparam,
+                           const std::vector<double>& aparam,
+                           const bool atomic);
+  void computew_mixed_type(std::vector<double>& ener,
+                           std::vector<float>& force,
+                           std::vector<float>& virial,
+                           std::vector<float>& atom_energy,
+                           std::vector<float>& atom_virial,
+                           const int& nframes,
+                           const std::vector<float>& coord,
+                           const std::vector<int>& atype,
+                           const std::vector<float>& box,
+                           const std::vector<float>& fparam,
+                           const std::vector<float>& aparam,
+                           const bool atomic);
 
  private:
   int num_intra_nthreads, num_inter_nthreads;
diff --git a/source/api_cc/include/DeepPotTF.h b/source/api_cc/include/DeepPotTF.h
index 699b0ff7fe..ffc3aab08b 100644
--- a/source/api_cc/include/DeepPotTF.h
+++ b/source/api_cc/include/DeepPotTF.h
@@ -60,6 +60,7 @@ class DeepPotTF : public DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Whether to compute atomic energy and virial.
    **/
   template <typename VALUETYPE, typename ENERGYVTYPE>
   void compute(ENERGYVTYPE& ener,
@@ -70,8 +71,9 @@ class DeepPotTF : public DeepPotBase {
                const std::vector<VALUETYPE>& coord,
                const std::vector<int>& atype,
                const std::vector<VALUETYPE>& box,
-               const std::vector<VALUETYPE>& fparam = std::vector<VALUETYPE>(),
-               const std::vector<VALUETYPE>& aparam = std::vector<VALUETYPE>());
+               const std::vector<VALUETYPE>& fparam,
+               const std::vector<VALUETYPE>& aparam,
+               const bool atomic);
   /**
    * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial
    *by using this DP.
@@ -96,6 +98,7 @@ class DeepPotTF : public DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Whether to compute atomic energy and virial.
    **/
   template <typename VALUETYPE, typename ENERGYVTYPE>
   void compute(ENERGYVTYPE& ener,
@@ -109,41 +112,9 @@ class DeepPotTF : public DeepPotBase {
                const int nghost,
                const InputNlist& lmp_list,
                const int& ago,
-               const std::vector<VALUETYPE>& fparam = std::vector<VALUETYPE>(),
-               const std::vector<VALUETYPE>& aparam = std::vector<VALUETYPE>());
-  /**
-   * @brief Evaluate the energy, force, and virial with the mixed type
-   *by using this DP.
-   * @param[out] ener The system energy.
-   * @param[out] force The force on each atom.
-   * @param[out] virial The virial.
-   * @param[in] nframes The number of frames.
-   * @param[in] coord The coordinates of atoms. The array should be of size
-   *nframes x natoms x 3.
-   * @param[in] atype The atom types. The array should be of size nframes x
-   *natoms.
-   * @param[in] box The cell of the region. The array should be of size nframes
-   *x 9.
-   * @param[in] fparam The frame parameter. The array can be of size :
-   * nframes x dim_fparam.
-   * dim_fparam. Then all frames are assumed to be provided with the same
-   *fparam.
-   * @param[in] aparam The atomic parameter The array can be of size :
-   * nframes x natoms x dim_aparam.
-   * natoms x dim_aparam. Then all frames are assumed to be provided with the
-   *same aparam.
-   **/
-  template <typename VALUETYPE, typename ENERGYVTYPE>
-  void compute_mixed_type(
-      ENERGYVTYPE& ener,
-      std::vector<VALUETYPE>& force,
-      std::vector<VALUETYPE>& virial,
-      const int& nframes,
-      const std::vector<VALUETYPE>& coord,
-      const std::vector<int>& atype,
-      const std::vector<VALUETYPE>& box,
-      const std::vector<VALUETYPE>& fparam = std::vector<VALUETYPE>(),
-      const std::vector<VALUETYPE>& aparam = std::vector<VALUETYPE>());
+               const std::vector<VALUETYPE>& fparam,
+               const std::vector<VALUETYPE>& aparam,
+               const bool atomic);
   /**
    * @brief Evaluate the energy, force, and virial with the mixed type
    *by using this DP.
@@ -167,20 +138,21 @@ class DeepPotTF : public DeepPotBase {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] atomic Whether to compute atomic energy and virial.
    **/
   template <typename VALUETYPE, typename ENERGYVTYPE>
-  void compute_mixed_type(
-      ENERGYVTYPE& ener,
-      std::vector<VALUETYPE>& force,
-      std::vector<VALUETYPE>& virial,
-      std::vector<VALUETYPE>& atom_energy,
-      std::vector<VALUETYPE>& atom_virial,
-      const int& nframes,
-      const std::vector<VALUETYPE>& coord,
-      const std::vector<int>& atype,
-      const std::vector<VALUETYPE>& box,
-      const std::vector<VALUETYPE>& fparam = std::vector<VALUETYPE>(),
-      const std::vector<VALUETYPE>& aparam = std::vector<VALUETYPE>());
+  void compute_mixed_type(ENERGYVTYPE& ener,
+                          std::vector<VALUETYPE>& force,
+                          std::vector<VALUETYPE>& virial,
+                          std::vector<VALUETYPE>& atom_energy,
+                          std::vector<VALUETYPE>& atom_virial,
+                          const int& nframes,
+                          const std::vector<VALUETYPE>& coord,
+                          const std::vector<int>& atype,
+                          const std::vector<VALUETYPE>& box,
+                          const std::vector<VALUETYPE>& fparam,
+                          const std::vector<VALUETYPE>& aparam,
+                          const bool atomic);
 
  public:
   /**
@@ -248,8 +220,9 @@ class DeepPotTF : public DeepPotBase {
                 const std::vector<double>& coord,
                 const std::vector<int>& atype,
                 const std::vector<double>& box,
-                const std::vector<double>& fparam = std::vector<double>(),
-                const std::vector<double>& aparam = std::vector<double>());
+                const std::vector<double>& fparam,
+                const std::vector<double>& aparam,
+                const bool atomic);
   void computew(std::vector<double>& ener,
                 std::vector<float>& force,
                 std::vector<float>& virial,
@@ -258,8 +231,9 @@ class DeepPotTF : public DeepPotBase {
                 const std::vector<float>& coord,
                 const std::vector<int>& atype,
                 const std::vector<float>& box,
-                const std::vector<float>& fparam = std::vector<float>(),
-                const std::vector<float>& aparam = std::vector<float>());
+                const std::vector<float>& fparam,
+                const std::vector<float>& aparam,
+                const bool atomic);
   void computew(std::vector<double>& ener,
                 std::vector<double>& force,
                 std::vector<double>& virial,
@@ -271,8 +245,9 @@ class DeepPotTF : public DeepPotBase {
                 const int nghost,
                 const InputNlist& inlist,
                 const int& ago,
-                const std::vector<double>& fparam = std::vector<double>(),
-                const std::vector<double>& aparam = std::vector<double>());
+                const std::vector<double>& fparam,
+                const std::vector<double>& aparam,
+                const bool atomic);
   void computew(std::vector<double>& ener,
                 std::vector<float>& force,
                 std::vector<float>& virial,
@@ -284,32 +259,33 @@ class DeepPotTF : public DeepPotBase {
                 const int nghost,
                 const InputNlist& inlist,
                 const int& ago,
-                const std::vector<float>& fparam = std::vector<float>(),
-                const std::vector<float>& aparam = std::vector<float>());
-  void computew_mixed_type(
-      std::vector<double>& ener,
-      std::vector<double>& force,
-      std::vector<double>& virial,
-      std::vector<double>& atom_energy,
-      std::vector<double>& atom_virial,
-      const int& nframes,
-      const std::vector<double>& coord,
-      const std::vector<int>& atype,
-      const std::vector<double>& box,
-      const std::vector<double>& fparam = std::vector<double>(),
-      const std::vector<double>& aparam = std::vector<double>());
-  void computew_mixed_type(
-      std::vector<double>& ener,
-      std::vector<float>& force,
-      std::vector<float>& virial,
-      std::vector<float>& atom_energy,
-      std::vector<float>& atom_virial,
-      const int& nframes,
-      const std::vector<float>& coord,
-      const std::vector<int>& atype,
-      const std::vector<float>& box,
-      const std::vector<float>& fparam = std::vector<float>(),
-      const std::vector<float>& aparam = std::vector<float>());
+                const std::vector<float>& fparam,
+                const std::vector<float>& aparam,
+                const bool atomic);
+  void computew_mixed_type(std::vector<double>& ener,
+                           std::vector<double>& force,
+                           std::vector<double>& virial,
+                           std::vector<double>& atom_energy,
+                           std::vector<double>& atom_virial,
+                           const int& nframes,
+                           const std::vector<double>& coord,
+                           const std::vector<int>& atype,
+                           const std::vector<double>& box,
+                           const std::vector<double>& fparam,
+                           const std::vector<double>& aparam,
+                           const bool atomic);
+  void computew_mixed_type(std::vector<double>& ener,
+                           std::vector<float>& force,
+                           std::vector<float>& virial,
+                           std::vector<float>& atom_energy,
+                           std::vector<float>& atom_virial,
+                           const int& nframes,
+                           const std::vector<float>& coord,
+                           const std::vector<int>& atype,
+                           const std::vector<float>& box,
+                           const std::vector<float>& fparam,
+                           const std::vector<float>& aparam,
+                           const bool atomic);
 
  private:
   tensorflow::Session* session;
diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc
index dd04e8e5e1..c184446288 100644
--- a/source/api_cc/src/DeepPot.cc
+++ b/source/api_cc/src/DeepPot.cc
@@ -80,7 +80,7 @@ void DeepPot::compute(ENERGYTYPE& dener,
   std::vector<ENERGYTYPE> dener_;
   std::vector<VALUETYPE> datom_energy_, datom_virial_;
   dp->computew(dener_, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_,
-               datype_, dbox, fparam_, aparam_);
+               datype_, dbox, fparam_, aparam_, false);
   dener = dener_[0];
 }
 
@@ -95,7 +95,7 @@ void DeepPot::compute(std::vector<ENERGYTYPE>& dener,
                       const std::vector<VALUETYPE>& aparam_) {
   std::vector<VALUETYPE> datom_energy_, datom_virial_;
   dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_,
-               datype_, dbox, fparam_, aparam_);
+               datype_, dbox, fparam_, aparam_, false);
 }
 
 template void DeepPot::compute<double>(ENERGYTYPE& dener,
@@ -149,7 +149,7 @@ void DeepPot::compute(ENERGYTYPE& dener,
   std::vector<ENERGYTYPE> dener_;
   std::vector<VALUETYPE> datom_energy_, datom_virial_;
   dp->computew(dener_, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_,
-               datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__);
+               datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, false);
   dener = dener_[0];
 }
 
@@ -167,7 +167,7 @@ void DeepPot::compute(std::vector<ENERGYTYPE>& dener,
                       const std::vector<VALUETYPE>& aparam__) {
   std::vector<VALUETYPE> datom_energy_, datom_virial_;
   dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_,
-               datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__);
+               datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, false);
 }
 
 template void DeepPot::compute<double>(ENERGYTYPE& dener,
@@ -231,7 +231,7 @@ void DeepPot::compute(ENERGYTYPE& dener,
                       const std::vector<VALUETYPE>& aparam_) {
   std::vector<ENERGYTYPE> dener_;
   dp->computew(dener_, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_,
-               datype_, dbox, fparam_, aparam_);
+               datype_, dbox, fparam_, aparam_, true);
   dener = dener_[0];
 }
 template <typename VALUETYPE>
@@ -246,7 +246,7 @@ void DeepPot::compute(std::vector<ENERGYTYPE>& dener,
                       const std::vector<VALUETYPE>& fparam_,
                       const std::vector<VALUETYPE>& aparam_) {
   dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_,
-               datype_, dbox, fparam_, aparam_);
+               datype_, dbox, fparam_, aparam_, true);
 }
 
 template void DeepPot::compute<double>(ENERGYTYPE& dener,
@@ -309,7 +309,7 @@ void DeepPot::compute(ENERGYTYPE& dener,
                       const std::vector<VALUETYPE>& aparam__) {
   std::vector<ENERGYTYPE> dener_;
   dp->computew(dener_, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_,
-               datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__);
+               datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, true);
   dener = dener_[0];
 }
 template <typename VALUETYPE>
@@ -327,7 +327,7 @@ void DeepPot::compute(std::vector<ENERGYTYPE>& dener,
                       const std::vector<VALUETYPE>& fparam_,
                       const std::vector<VALUETYPE>& aparam__) {
   dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_,
-               datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__);
+               datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, true);
 }
 
 template void DeepPot::compute<double>(ENERGYTYPE& dener,
@@ -401,7 +401,7 @@ void DeepPot::compute_mixed_type(ENERGYTYPE& dener,
   std::vector<VALUETYPE> datom_energy_, datom_virial_;
   dp->computew_mixed_type(dener_, dforce_, dvirial, datom_energy_,
                           datom_virial_, nframes, dcoord_, datype_, dbox,
-                          fparam_, aparam_);
+                          fparam_, aparam_, false);
   dener = dener_[0];
 }
 template <typename VALUETYPE>
@@ -416,7 +416,8 @@ void DeepPot::compute_mixed_type(std::vector<ENERGYTYPE>& dener,
                                  const std::vector<VALUETYPE>& aparam_) {
   std::vector<VALUETYPE> datom_energy_, datom_virial_;
   dp->computew_mixed_type(dener, dforce_, dvirial, datom_energy_, datom_virial_,
-                          nframes, dcoord_, datype_, dbox, fparam_, aparam_);
+                          nframes, dcoord_, datype_, dbox, fparam_, aparam_,
+                          false);
 }
 
 template void DeepPot::compute_mixed_type<double>(
@@ -478,7 +479,7 @@ void DeepPot::compute_mixed_type(ENERGYTYPE& dener,
   std::vector<ENERGYTYPE> dener_;
   dp->computew_mixed_type(dener_, dforce_, dvirial, datom_energy_,
                           datom_virial_, nframes, dcoord_, datype_, dbox,
-                          fparam_, aparam_);
+                          fparam_, aparam_, true);
   dener = dener_[0];
 }
 template <typename VALUETYPE>
@@ -494,7 +495,8 @@ void DeepPot::compute_mixed_type(std::vector<ENERGYTYPE>& dener,
                                  const std::vector<VALUETYPE>& fparam_,
                                  const std::vector<VALUETYPE>& aparam_) {
   dp->computew_mixed_type(dener, dforce_, dvirial, datom_energy_, datom_virial_,
-                          nframes, dcoord_, datype_, dbox, fparam_, aparam_);
+                          nframes, dcoord_, datype_, dbox, fparam_, aparam_,
+                          true);
 }
 
 template void DeepPot::compute_mixed_type<double>(
diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc
index 33597df4e0..0b4f87c150 100644
--- a/source/api_cc/src/DeepPotPT.cc
+++ b/source/api_cc/src/DeepPotPT.cc
@@ -125,7 +125,8 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
                         const InputNlist& lmp_list,
                         const int& ago,
                         const std::vector<VALUETYPE>& fparam,
-                        const std::vector<VALUETYPE>& aparam) {
+                        const std::vector<VALUETYPE>& aparam,
+                        const bool atomic) {
   torch::Device device(torch::kCUDA, gpu_id);
   if (!gpu_enabled) {
     device = torch::Device(torch::kCPU);
@@ -212,7 +213,7 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
   }
   at::Tensor firstneigh = createNlistTensor(nlist_data.jlist);
   firstneigh_tensor = firstneigh.to(torch::kInt64).to(device);
-  bool do_atom_virial_tensor = true;
+  bool do_atom_virial_tensor = atomic;
   c10::optional<torch::Tensor> optional_tensor;
   c10::optional<torch::Tensor> fparam_tensor;
   if (!fparam.empty()) {
@@ -246,20 +247,10 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
   c10::IValue energy_ = outputs.at("energy");
   c10::IValue force_ = outputs.at("extended_force");
   c10::IValue virial_ = outputs.at("virial");
-  c10::IValue atom_virial_ = outputs.at("extended_virial");
-  c10::IValue atom_energy_ = outputs.at("atom_energy");
   torch::Tensor flat_energy_ = energy_.toTensor().view({-1});
   torch::Tensor cpu_energy_ = flat_energy_.to(torch::kCPU);
   ener.assign(cpu_energy_.data_ptr<ENERGYTYPE>(),
               cpu_energy_.data_ptr<ENERGYTYPE>() + cpu_energy_.numel());
-  torch::Tensor flat_atom_energy_ =
-      atom_energy_.toTensor().view({-1}).to(floatType);
-  torch::Tensor cpu_atom_energy_ = flat_atom_energy_.to(torch::kCPU);
-  datom_energy.resize(nall_real,
-                      0.0);  // resize to nall to be consistenet with TF.
-  datom_energy.assign(
-      cpu_atom_energy_.data_ptr<VALUETYPE>(),
-      cpu_atom_energy_.data_ptr<VALUETYPE>() + cpu_atom_energy_.numel());
   torch::Tensor flat_force_ = force_.toTensor().view({-1}).to(floatType);
   torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU);
   dforce.assign(cpu_force_.data_ptr<VALUETYPE>(),
@@ -268,22 +259,35 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
   torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU);
   virial.assign(cpu_virial_.data_ptr<VALUETYPE>(),
                 cpu_virial_.data_ptr<VALUETYPE>() + cpu_virial_.numel());
-  torch::Tensor flat_atom_virial_ =
-      atom_virial_.toTensor().view({-1}).to(floatType);
-  torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU);
-  datom_virial.assign(
-      cpu_atom_virial_.data_ptr<VALUETYPE>(),
-      cpu_atom_virial_.data_ptr<VALUETYPE>() + cpu_atom_virial_.numel());
+
   // bkw map
   force.resize(static_cast<size_t>(nframes) * fwd_map.size() * 3);
-  atom_energy.resize(static_cast<size_t>(nframes) * fwd_map.size());
-  atom_virial.resize(static_cast<size_t>(nframes) * fwd_map.size() * 9);
   select_map<VALUETYPE>(force, dforce, bkw_map, 3, nframes, fwd_map.size(),
                         nall_real);
-  select_map<VALUETYPE>(atom_energy, datom_energy, bkw_map, 1, nframes,
-                        fwd_map.size(), nall_real);
-  select_map<VALUETYPE>(atom_virial, datom_virial, bkw_map, 9, nframes,
-                        fwd_map.size(), nall_real);
+  if (atomic) {
+    c10::IValue atom_virial_ = outputs.at("extended_virial");
+    c10::IValue atom_energy_ = outputs.at("atom_energy");
+    torch::Tensor flat_atom_energy_ =
+        atom_energy_.toTensor().view({-1}).to(floatType);
+    torch::Tensor cpu_atom_energy_ = flat_atom_energy_.to(torch::kCPU);
+    datom_energy.resize(nall_real,
+                        0.0);  // resize to nall to be consistenet with TF.
+    datom_energy.assign(
+        cpu_atom_energy_.data_ptr<VALUETYPE>(),
+        cpu_atom_energy_.data_ptr<VALUETYPE>() + cpu_atom_energy_.numel());
+    torch::Tensor flat_atom_virial_ =
+        atom_virial_.toTensor().view({-1}).to(floatType);
+    torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU);
+    datom_virial.assign(
+        cpu_atom_virial_.data_ptr<VALUETYPE>(),
+        cpu_atom_virial_.data_ptr<VALUETYPE>() + cpu_atom_virial_.numel());
+    atom_energy.resize(static_cast<size_t>(nframes) * fwd_map.size());
+    atom_virial.resize(static_cast<size_t>(nframes) * fwd_map.size() * 9);
+    select_map<VALUETYPE>(atom_energy, datom_energy, bkw_map, 1, nframes,
+                          fwd_map.size(), nall_real);
+    select_map<VALUETYPE>(atom_virial, datom_virial, bkw_map, 9, nframes,
+                          fwd_map.size(), nall_real);
+  }
 }
 template void DeepPotPT::compute<double, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& ener,
@@ -298,7 +302,8 @@ template void DeepPotPT::compute<double, std::vector<ENERGYTYPE>>(
     const InputNlist& lmp_list,
     const int& ago,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam);
+    const std::vector<double>& aparam,
+    const bool atomic);
 template void DeepPotPT::compute<float, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& ener,
     std::vector<float>& force,
@@ -312,7 +317,8 @@ template void DeepPotPT::compute<float, std::vector<ENERGYTYPE>>(
     const InputNlist& lmp_list,
     const int& ago,
     const std::vector<float>& fparam,
-    const std::vector<float>& aparam);
+    const std::vector<float>& aparam,
+    const bool atomic);
 template <typename VALUETYPE, typename ENERGYVTYPE>
 void DeepPotPT::compute(ENERGYVTYPE& ener,
                         std::vector<VALUETYPE>& force,
@@ -323,7 +329,8 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
                         const std::vector<int>& atype,
                         const std::vector<VALUETYPE>& box,
                         const std::vector<VALUETYPE>& fparam,
-                        const std::vector<VALUETYPE>& aparam) {
+                        const std::vector<VALUETYPE>& aparam,
+                        const bool atomic) {
   torch::Device device(torch::kCUDA, gpu_id);
   if (!gpu_enabled) {
     device = torch::Device(torch::kCPU);
@@ -389,25 +396,17 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
             .to(device);
   }
   inputs.push_back(aparam_tensor);
-  bool do_atom_virial_tensor = true;
+  bool do_atom_virial_tensor = atomic;
   inputs.push_back(do_atom_virial_tensor);
   c10::Dict<c10::IValue, c10::IValue> outputs =
       module.forward(inputs).toGenericDict();
   c10::IValue energy_ = outputs.at("energy");
   c10::IValue force_ = outputs.at("force");
   c10::IValue virial_ = outputs.at("virial");
-  c10::IValue atom_virial_ = outputs.at("atom_virial");
-  c10::IValue atom_energy_ = outputs.at("atom_energy");
   torch::Tensor flat_energy_ = energy_.toTensor().view({-1});
   torch::Tensor cpu_energy_ = flat_energy_.to(torch::kCPU);
   ener.assign(cpu_energy_.data_ptr<ENERGYTYPE>(),
               cpu_energy_.data_ptr<ENERGYTYPE>() + cpu_energy_.numel());
-  torch::Tensor flat_atom_energy_ =
-      atom_energy_.toTensor().view({-1}).to(floatType);
-  torch::Tensor cpu_atom_energy_ = flat_atom_energy_.to(torch::kCPU);
-  atom_energy.assign(
-      cpu_atom_energy_.data_ptr<VALUETYPE>(),
-      cpu_atom_energy_.data_ptr<VALUETYPE>() + cpu_atom_energy_.numel());
   torch::Tensor flat_force_ = force_.toTensor().view({-1}).to(floatType);
   torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU);
   force.assign(cpu_force_.data_ptr<VALUETYPE>(),
@@ -416,12 +415,22 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
   torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU);
   virial.assign(cpu_virial_.data_ptr<VALUETYPE>(),
                 cpu_virial_.data_ptr<VALUETYPE>() + cpu_virial_.numel());
-  torch::Tensor flat_atom_virial_ =
-      atom_virial_.toTensor().view({-1}).to(floatType);
-  torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU);
-  atom_virial.assign(
-      cpu_atom_virial_.data_ptr<VALUETYPE>(),
-      cpu_atom_virial_.data_ptr<VALUETYPE>() + cpu_atom_virial_.numel());
+  if (atomic) {
+    c10::IValue atom_virial_ = outputs.at("atom_virial");
+    c10::IValue atom_energy_ = outputs.at("atom_energy");
+    torch::Tensor flat_atom_energy_ =
+        atom_energy_.toTensor().view({-1}).to(floatType);
+    torch::Tensor cpu_atom_energy_ = flat_atom_energy_.to(torch::kCPU);
+    atom_energy.assign(
+        cpu_atom_energy_.data_ptr<VALUETYPE>(),
+        cpu_atom_energy_.data_ptr<VALUETYPE>() + cpu_atom_energy_.numel());
+    torch::Tensor flat_atom_virial_ =
+        atom_virial_.toTensor().view({-1}).to(floatType);
+    torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU);
+    atom_virial.assign(
+        cpu_atom_virial_.data_ptr<VALUETYPE>(),
+        cpu_atom_virial_.data_ptr<VALUETYPE>() + cpu_atom_virial_.numel());
+  }
 }
 
 template void DeepPotPT::compute<double, std::vector<ENERGYTYPE>>(
@@ -434,7 +443,8 @@ template void DeepPotPT::compute<double, std::vector<ENERGYTYPE>>(
     const std::vector<int>& atype,
     const std::vector<double>& box,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam);
+    const std::vector<double>& aparam,
+    const bool atomic);
 template void DeepPotPT::compute<float, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& ener,
     std::vector<float>& force,
@@ -445,7 +455,8 @@ template void DeepPotPT::compute<float, std::vector<ENERGYTYPE>>(
     const std::vector<int>& atype,
     const std::vector<float>& box,
     const std::vector<float>& fparam,
-    const std::vector<float>& aparam);
+    const std::vector<float>& aparam,
+    const bool atomic);
 void DeepPotPT::get_type_map(std::string& type_map) {
   auto ret = module.run_method("get_type_map").toList();
   for (const torch::IValue& element : ret) {
@@ -464,10 +475,11 @@ void DeepPotPT::computew(std::vector<double>& ener,
                          const std::vector<int>& atype,
                          const std::vector<double>& box,
                          const std::vector<double>& fparam,
-                         const std::vector<double>& aparam) {
+                         const std::vector<double>& aparam,
+                         const bool atomic) {
   translate_error([&] {
     compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box,
-            fparam, aparam);
+            fparam, aparam, atomic);
   });
 }
 void DeepPotPT::computew(std::vector<double>& ener,
@@ -479,10 +491,11 @@ void DeepPotPT::computew(std::vector<double>& ener,
                          const std::vector<int>& atype,
                          const std::vector<float>& box,
                          const std::vector<float>& fparam,
-                         const std::vector<float>& aparam) {
+                         const std::vector<float>& aparam,
+                         const bool atomic) {
   translate_error([&] {
     compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box,
-            fparam, aparam);
+            fparam, aparam, atomic);
   });
 }
 void DeepPotPT::computew(std::vector<double>& ener,
@@ -497,10 +510,11 @@ void DeepPotPT::computew(std::vector<double>& ener,
                          const InputNlist& inlist,
                          const int& ago,
                          const std::vector<double>& fparam,
-                         const std::vector<double>& aparam) {
+                         const std::vector<double>& aparam,
+                         const bool atomic) {
   translate_error([&] {
     compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box,
-            nghost, inlist, ago, fparam, aparam);
+            nghost, inlist, ago, fparam, aparam, atomic);
   });
 }
 void DeepPotPT::computew(std::vector<double>& ener,
@@ -515,10 +529,11 @@ void DeepPotPT::computew(std::vector<double>& ener,
                          const InputNlist& inlist,
                          const int& ago,
                          const std::vector<float>& fparam,
-                         const std::vector<float>& aparam) {
+                         const std::vector<float>& aparam,
+                         const bool atomic) {
   translate_error([&] {
     compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box,
-            nghost, inlist, ago, fparam, aparam);
+            nghost, inlist, ago, fparam, aparam, atomic);
   });
 }
 void DeepPotPT::computew_mixed_type(std::vector<double>& ener,
@@ -531,7 +546,8 @@ void DeepPotPT::computew_mixed_type(std::vector<double>& ener,
                                     const std::vector<int>& atype,
                                     const std::vector<double>& box,
                                     const std::vector<double>& fparam,
-                                    const std::vector<double>& aparam) {
+                                    const std::vector<double>& aparam,
+                                    const bool atomic) {
   throw deepmd::deepmd_exception("computew_mixed_type is not implemented");
 }
 void DeepPotPT::computew_mixed_type(std::vector<double>& ener,
@@ -544,7 +560,8 @@ void DeepPotPT::computew_mixed_type(std::vector<double>& ener,
                                     const std::vector<int>& atype,
                                     const std::vector<float>& box,
                                     const std::vector<float>& fparam,
-                                    const std::vector<float>& aparam) {
+                                    const std::vector<float>& aparam,
+                                    const bool atomic) {
   throw deepmd::deepmd_exception("computew_mixed_type is not implemented");
 }
 #endif
diff --git a/source/api_cc/src/DeepPotTF.cc b/source/api_cc/src/DeepPotTF.cc
index 6b39e0ada9..2c09c17a69 100644
--- a/source/api_cc/src/DeepPotTF.cc
+++ b/source/api_cc/src/DeepPotTF.cc
@@ -581,7 +581,8 @@ void DeepPotTF::compute(ENERGYVTYPE& dener,
                         const std::vector<int>& datype_,
                         const std::vector<VALUETYPE>& dbox,
                         const std::vector<VALUETYPE>& fparam_,
-                        const std::vector<VALUETYPE>& aparam_) {
+                        const std::vector<VALUETYPE>& aparam_,
+                        const bool atomic) {
   // if datype.size is 0, not clear nframes; but 1 is just ok
   int nframes = datype_.size() > 0 ? (dcoord_.size() / 3 / datype_.size()) : 1;
   atommap = deepmd::AtomMap(datype_.begin(), datype_.end());
@@ -598,14 +599,24 @@ void DeepPotTF::compute(ENERGYVTYPE& dener,
     int ret = session_input_tensors<double>(input_tensors, dcoord_, ntypes,
                                             datype_, dbox, cell_size, fparam,
                                             aparam, atommap, "", aparam_nall);
-    run_model<double>(dener, dforce_, dvirial, datom_energy_, datom_virial_,
-                      session, input_tensors, atommap, nframes);
+    if (atomic) {
+      run_model<double>(dener, dforce_, dvirial, datom_energy_, datom_virial_,
+                        session, input_tensors, atommap, nframes);
+    } else {
+      run_model<double>(dener, dforce_, dvirial, session, input_tensors,
+                        atommap, nframes);
+    }
   } else {
     int ret = session_input_tensors<float>(input_tensors, dcoord_, ntypes,
                                            datype_, dbox, cell_size, fparam,
                                            aparam, atommap, "", aparam_nall);
-    run_model<float>(dener, dforce_, dvirial, datom_energy_, datom_virial_,
-                     session, input_tensors, atommap, nframes);
+    if (atomic) {
+      run_model<float>(dener, dforce_, dvirial, datom_energy_, datom_virial_,
+                       session, input_tensors, atommap, nframes);
+    } else {
+      run_model<float>(dener, dforce_, dvirial, session, input_tensors, atommap,
+                       nframes);
+    }
   }
 }
 
@@ -619,7 +630,8 @@ template void DeepPotTF::compute<double, ENERGYTYPE>(
     const std::vector<int>& datype_,
     const std::vector<double>& dbox,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam);
+    const std::vector<double>& aparam,
+    const bool atomic);
 
 template void DeepPotTF::compute<float, ENERGYTYPE>(
     ENERGYTYPE& dener,
@@ -631,7 +643,8 @@ template void DeepPotTF::compute<float, ENERGYTYPE>(
     const std::vector<int>& datype_,
     const std::vector<float>& dbox,
     const std::vector<float>& fparam,
-    const std::vector<float>& aparam);
+    const std::vector<float>& aparam,
+    const bool atomic);
 
 template void DeepPotTF::compute<double, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& dener,
@@ -643,7 +656,8 @@ template void DeepPotTF::compute<double, std::vector<ENERGYTYPE>>(
     const std::vector<int>& datype_,
     const std::vector<double>& dbox,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam);
+    const std::vector<double>& aparam,
+    const bool atomic);
 
 template void DeepPotTF::compute<float, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& dener,
@@ -655,7 +669,8 @@ template void DeepPotTF::compute<float, std::vector<ENERGYTYPE>>(
     const std::vector<int>& datype_,
     const std::vector<float>& dbox,
     const std::vector<float>& fparam,
-    const std::vector<float>& aparam);
+    const std::vector<float>& aparam,
+    const bool atomic);
 
 template <typename VALUETYPE, typename ENERGYVTYPE>
 void DeepPotTF::compute(ENERGYVTYPE& dener,
@@ -670,7 +685,8 @@ void DeepPotTF::compute(ENERGYVTYPE& dener,
                         const InputNlist& lmp_list,
                         const int& ago,
                         const std::vector<VALUETYPE>& fparam_,
-                        const std::vector<VALUETYPE>& aparam__) {
+                        const std::vector<VALUETYPE>& aparam__,
+                        const bool atomic) {
   int nall = datype_.size();
   // if nall==0, unclear nframes, but 1 is ok
   int nframes = nall > 0 ? (dcoord_.size() / nall / 3) : 1;
@@ -706,15 +722,25 @@ void DeepPotTF::compute(ENERGYVTYPE& dener,
         input_tensors, dcoord, ntypes, datype, dbox, nlist, fparam, aparam,
         atommap, nghost_real, ago, "", aparam_nall);
     assert(nloc_real == ret);
-    run_model<double>(dener, dforce, dvirial, datom_energy, datom_virial,
-                      session, input_tensors, atommap, nframes, nghost_real);
+    if (atomic) {
+      run_model<double>(dener, dforce, dvirial, datom_energy, datom_virial,
+                        session, input_tensors, atommap, nframes, nghost_real);
+    } else {
+      run_model<double>(dener, dforce, dvirial, session, input_tensors, atommap,
+                        nframes, nghost_real);
+    }
   } else {
     int ret = session_input_tensors<float>(
         input_tensors, dcoord, ntypes, datype, dbox, nlist, fparam, aparam,
         atommap, nghost_real, ago, "", aparam_nall);
     assert(nloc_real == ret);
-    run_model<float>(dener, dforce, dvirial, datom_energy, datom_virial,
-                     session, input_tensors, atommap, nframes, nghost_real);
+    if (atomic) {
+      run_model<float>(dener, dforce, dvirial, datom_energy, datom_virial,
+                       session, input_tensors, atommap, nframes, nghost_real);
+    } else {
+      run_model<float>(dener, dforce, dvirial, session, input_tensors, atommap,
+                       nframes, nghost_real);
+    }
   }
 
   // bkw map
@@ -742,7 +768,8 @@ template void DeepPotTF::compute<double, ENERGYTYPE>(
     const InputNlist& lmp_list,
     const int& ago,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam_);
+    const std::vector<double>& aparam_,
+    const bool atomic);
 
 template void DeepPotTF::compute<float, ENERGYTYPE>(
     ENERGYTYPE& dener,
@@ -757,7 +784,8 @@ template void DeepPotTF::compute<float, ENERGYTYPE>(
     const InputNlist& lmp_list,
     const int& ago,
     const std::vector<float>& fparam,
-    const std::vector<float>& aparam_);
+    const std::vector<float>& aparam_,
+    const bool atomic);
 
 template void DeepPotTF::compute<double, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& dener,
@@ -772,7 +800,8 @@ template void DeepPotTF::compute<double, std::vector<ENERGYTYPE>>(
     const InputNlist& lmp_list,
     const int& ago,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam_);
+    const std::vector<double>& aparam_,
+    const bool atomic);
 
 template void DeepPotTF::compute<float, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& dener,
@@ -787,90 +816,10 @@ template void DeepPotTF::compute<float, std::vector<ENERGYTYPE>>(
     const InputNlist& lmp_list,
     const int& ago,
     const std::vector<float>& fparam,
-    const std::vector<float>& aparam_);
+    const std::vector<float>& aparam_,
+    const bool atomic);
 
 // mixed type
-template <typename VALUETYPE, typename ENERGYVTYPE>
-void DeepPotTF::compute_mixed_type(ENERGYVTYPE& dener,
-                                   std::vector<VALUETYPE>& dforce_,
-                                   std::vector<VALUETYPE>& dvirial,
-                                   const int& nframes,
-                                   const std::vector<VALUETYPE>& dcoord_,
-                                   const std::vector<int>& datype_,
-                                   const std::vector<VALUETYPE>& dbox,
-                                   const std::vector<VALUETYPE>& fparam_,
-                                   const std::vector<VALUETYPE>& aparam_) {
-  int nloc = datype_.size() / nframes;
-  // here atommap only used to get nloc
-  atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc);
-  std::vector<VALUETYPE> fparam;
-  std::vector<VALUETYPE> aparam;
-  validate_fparam_aparam(nframes, nloc, fparam_, aparam_);
-  tile_fparam_aparam(fparam, nframes, dfparam, fparam_);
-  tile_fparam_aparam(aparam, nframes, nloc * daparam, aparam_);
-
-  std::vector<std::pair<std::string, Tensor>> input_tensors;
-
-  if (dtype == tensorflow::DT_DOUBLE) {
-    int ret = session_input_tensors_mixed_type<double>(
-        input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size,
-        fparam, aparam, atommap, "", aparam_nall);
-    assert(ret == nloc);
-    run_model<double>(dener, dforce_, dvirial, session, input_tensors, atommap,
-                      nframes);
-  } else {
-    int ret = session_input_tensors_mixed_type<float>(
-        input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size,
-        fparam, aparam, atommap, "", aparam_nall);
-    assert(ret == nloc);
-    run_model<float>(dener, dforce_, dvirial, session, input_tensors, atommap,
-                     nframes);
-  }
-}
-
-template void DeepPotTF::compute_mixed_type<double, ENERGYTYPE>(
-    ENERGYTYPE& dener,
-    std::vector<double>& dforce_,
-    std::vector<double>& dvirial,
-    const int& nframes,
-    const std::vector<double>& dcoord_,
-    const std::vector<int>& datype_,
-    const std::vector<double>& dbox,
-    const std::vector<double>& fparam,
-    const std::vector<double>& aparam);
-
-template void DeepPotTF::compute_mixed_type<float, ENERGYTYPE>(
-    ENERGYTYPE& dener,
-    std::vector<float>& dforce_,
-    std::vector<float>& dvirial,
-    const int& nframes,
-    const std::vector<float>& dcoord_,
-    const std::vector<int>& datype_,
-    const std::vector<float>& dbox,
-    const std::vector<float>& fparam,
-    const std::vector<float>& aparam);
-
-template void DeepPotTF::compute_mixed_type<double, std::vector<ENERGYTYPE>>(
-    std::vector<ENERGYTYPE>& dener,
-    std::vector<double>& dforce_,
-    std::vector<double>& dvirial,
-    const int& nframes,
-    const std::vector<double>& dcoord_,
-    const std::vector<int>& datype_,
-    const std::vector<double>& dbox,
-    const std::vector<double>& fparam,
-    const std::vector<double>& aparam);
-
-template void DeepPotTF::compute_mixed_type<float, std::vector<ENERGYTYPE>>(
-    std::vector<ENERGYTYPE>& dener,
-    std::vector<float>& dforce_,
-    std::vector<float>& dvirial,
-    const int& nframes,
-    const std::vector<float>& dcoord_,
-    const std::vector<int>& datype_,
-    const std::vector<float>& dbox,
-    const std::vector<float>& fparam,
-    const std::vector<float>& aparam);
 
 template <typename VALUETYPE, typename ENERGYVTYPE>
 void DeepPotTF::compute_mixed_type(ENERGYVTYPE& dener,
@@ -883,7 +832,8 @@ void DeepPotTF::compute_mixed_type(ENERGYVTYPE& dener,
                                    const std::vector<int>& datype_,
                                    const std::vector<VALUETYPE>& dbox,
                                    const std::vector<VALUETYPE>& fparam_,
-                                   const std::vector<VALUETYPE>& aparam_) {
+                                   const std::vector<VALUETYPE>& aparam_,
+                                   const bool atomic) {
   int nloc = datype_.size() / nframes;
   // here atommap only used to get nloc
   atommap = deepmd::AtomMap(datype_.begin(), datype_.begin() + nloc);
@@ -899,14 +849,24 @@ void DeepPotTF::compute_mixed_type(ENERGYVTYPE& dener,
     int nloc = session_input_tensors_mixed_type<double>(
         input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size,
         fparam, aparam, atommap, "", aparam_nall);
-    run_model<double>(dener, dforce_, dvirial, datom_energy_, datom_virial_,
-                      session, input_tensors, atommap, nframes);
+    if (atomic) {
+      run_model<double>(dener, dforce_, dvirial, datom_energy_, datom_virial_,
+                        session, input_tensors, atommap, nframes);
+    } else {
+      run_model<double>(dener, dforce_, dvirial, session, input_tensors,
+                        atommap, nframes);
+    }
   } else {
     int nloc = session_input_tensors_mixed_type<float>(
         input_tensors, nframes, dcoord_, ntypes, datype_, dbox, cell_size,
         fparam, aparam, atommap, "", aparam_nall);
-    run_model<float>(dener, dforce_, dvirial, datom_energy_, datom_virial_,
-                     session, input_tensors, atommap, nframes);
+    if (atomic) {
+      run_model<float>(dener, dforce_, dvirial, datom_energy_, datom_virial_,
+                       session, input_tensors, atommap, nframes);
+    } else {
+      run_model<float>(dener, dforce_, dvirial, session, input_tensors, atommap,
+                       nframes);
+    }
   }
 }
 
@@ -921,7 +881,8 @@ template void DeepPotTF::compute_mixed_type<double, ENERGYTYPE>(
     const std::vector<int>& datype_,
     const std::vector<double>& dbox,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam);
+    const std::vector<double>& aparam,
+    const bool atomic);
 
 template void DeepPotTF::compute_mixed_type<float, ENERGYTYPE>(
     ENERGYTYPE& dener,
@@ -934,7 +895,8 @@ template void DeepPotTF::compute_mixed_type<float, ENERGYTYPE>(
     const std::vector<int>& datype_,
     const std::vector<float>& dbox,
     const std::vector<float>& fparam,
-    const std::vector<float>& aparam);
+    const std::vector<float>& aparam,
+    const bool atomic);
 
 template void DeepPotTF::compute_mixed_type<double, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& dener,
@@ -947,7 +909,8 @@ template void DeepPotTF::compute_mixed_type<double, std::vector<ENERGYTYPE>>(
     const std::vector<int>& datype_,
     const std::vector<double>& dbox,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam);
+    const std::vector<double>& aparam,
+    const bool atomic);
 
 template void DeepPotTF::compute_mixed_type<float, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& dener,
@@ -960,7 +923,8 @@ template void DeepPotTF::compute_mixed_type<float, std::vector<ENERGYTYPE>>(
     const std::vector<int>& datype_,
     const std::vector<float>& dbox,
     const std::vector<float>& fparam,
-    const std::vector<float>& aparam);
+    const std::vector<float>& aparam,
+    const bool atomic);
 
 void DeepPotTF::get_type_map(std::string& type_map) {
   type_map = get_scalar<STRINGTYPE>("model_attr/tmap");
@@ -976,9 +940,10 @@ void DeepPotTF::computew(std::vector<double>& ener,
                          const std::vector<int>& atype,
                          const std::vector<double>& box,
                          const std::vector<double>& fparam,
-                         const std::vector<double>& aparam) {
+                         const std::vector<double>& aparam,
+                         const bool atomic) {
   compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box,
-          fparam, aparam);
+          fparam, aparam, atomic);
 }
 void DeepPotTF::computew(std::vector<double>& ener,
                          std::vector<float>& force,
@@ -989,9 +954,10 @@ void DeepPotTF::computew(std::vector<double>& ener,
                          const std::vector<int>& atype,
                          const std::vector<float>& box,
                          const std::vector<float>& fparam,
-                         const std::vector<float>& aparam) {
+                         const std::vector<float>& aparam,
+                         const bool atomic) {
   compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box,
-          fparam, aparam);
+          fparam, aparam, atomic);
 }
 void DeepPotTF::computew(std::vector<double>& ener,
                          std::vector<double>& force,
@@ -1005,9 +971,10 @@ void DeepPotTF::computew(std::vector<double>& ener,
                          const InputNlist& inlist,
                          const int& ago,
                          const std::vector<double>& fparam,
-                         const std::vector<double>& aparam) {
+                         const std::vector<double>& aparam,
+                         const bool atomic) {
   compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box,
-          nghost, inlist, ago, fparam, aparam);
+          nghost, inlist, ago, fparam, aparam, atomic);
 }
 void DeepPotTF::computew(std::vector<double>& ener,
                          std::vector<float>& force,
@@ -1021,9 +988,10 @@ void DeepPotTF::computew(std::vector<double>& ener,
                          const InputNlist& inlist,
                          const int& ago,
                          const std::vector<float>& fparam,
-                         const std::vector<float>& aparam) {
+                         const std::vector<float>& aparam,
+                         const bool atomic) {
   compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box,
-          nghost, inlist, ago, fparam, aparam);
+          nghost, inlist, ago, fparam, aparam, atomic);
 }
 void DeepPotTF::computew_mixed_type(std::vector<double>& ener,
                                     std::vector<double>& force,
@@ -1035,9 +1003,10 @@ void DeepPotTF::computew_mixed_type(std::vector<double>& ener,
                                     const std::vector<int>& atype,
                                     const std::vector<double>& box,
                                     const std::vector<double>& fparam,
-                                    const std::vector<double>& aparam) {
+                                    const std::vector<double>& aparam,
+                                    const bool atomic) {
   compute_mixed_type(ener, force, virial, atom_energy, atom_virial, nframes,
-                     coord, atype, box, fparam, aparam);
+                     coord, atype, box, fparam, aparam, atomic);
 }
 void DeepPotTF::computew_mixed_type(std::vector<double>& ener,
                                     std::vector<float>& force,
@@ -1049,8 +1018,9 @@ void DeepPotTF::computew_mixed_type(std::vector<double>& ener,
                                     const std::vector<int>& atype,
                                     const std::vector<float>& box,
                                     const std::vector<float>& fparam,
-                                    const std::vector<float>& aparam) {
+                                    const std::vector<float>& aparam,
+                                    const bool atomic) {
   compute_mixed_type(ener, force, virial, atom_energy, atom_virial, nframes,
-                     coord, atype, box, fparam, aparam);
+                     coord, atype, box, fparam, aparam, atomic);
 }
 #endif