diff --git a/source/lib/include/common.h b/source/lib/include/common.h
index 3912f21f7f..620c88eb6f 100644
--- a/source/lib/include/common.h
+++ b/source/lib/include/common.h
@@ -1,14 +1,22 @@
 #pragma once
 
 #include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/version.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/shape_inference.h"
+#include <string>
 
 using namespace tensorflow;
 using namespace std;
 
+#if TF_MAJOR_VERSION >= 2 && TF_MINOR_VERSION >= 2
+typedef tensorflow::tstring STRINGTYPE;
+#else
+typedef std::string STRINGTYPE;
+#endif
+
 #include "NNPAtomMap.h"
 #include <vector>
 #include "version.h"
diff --git a/source/lib/src/DataModifier.cc b/source/lib/src/DataModifier.cc
index 65b3f85dd4..03fcd43a07 100644
--- a/source/lib/src/DataModifier.cc
+++ b/source/lib/src/DataModifier.cc
@@ -37,7 +37,7 @@ init (const string & model,
   rcut = get_scalar<VALUETYPE>("descrpt_attr/rcut");
   cell_size = rcut;
   ntypes = get_scalar<int>("descrpt_attr/ntypes");
-  model_type = get_scalar<string>("model_attr/model_type");
+  model_type = get_scalar<STRINGTYPE>("model_attr/model_type");
   get_vector<int>(sel_type, "model_attr/sel_type");
   sort(sel_type.begin(), sel_type.end());
   inited = true;
diff --git a/source/lib/src/DeepTensor.cc b/source/lib/src/DeepTensor.cc
index bd3a0ea138..9b9715fd0d 100644
--- a/source/lib/src/DeepTensor.cc
+++ b/source/lib/src/DeepTensor.cc
@@ -33,7 +33,7 @@ init (const string & model,
   rcut = get_scalar<VALUETYPE>("descrpt_attr/rcut");
   cell_size = rcut;
   ntypes = get_scalar<int>("descrpt_attr/ntypes");
-  model_type = get_scalar<string>("model_attr/model_type");
+  model_type = get_scalar<STRINGTYPE>("model_attr/model_type");
   odim = get_scalar<int>("model_attr/output_dim");
   get_vector<int>(sel_type, "model_attr/sel_type");
   inited = true;
diff --git a/source/lib/src/NNPInter.cc b/source/lib/src/NNPInter.cc
index aea9d48c9b..c1347504bb 100644
--- a/source/lib/src/NNPInter.cc
+++ b/source/lib/src/NNPInter.cc
@@ -639,7 +639,7 @@ compute (ENERGYTYPE &			dener,
 void
 NNPInter::
 get_type_map(std::string & type_map){
-    type_map = get_scalar<std::string>("model_attr/tmap");
+    type_map = get_scalar<STRINGTYPE>("model_attr/tmap");
 }
 
 
diff --git a/source/train/Data.py b/source/train/Data.py
index 149c76bbe8..b38839a546 100644
--- a/source/train/Data.py
+++ b/source/train/Data.py
@@ -476,7 +476,7 @@ def load_energy(self,
         # ignore energy_file
         if coeff_atom_ener == 1:
             ener = np.sum(atom_ener, axis = 1)
-            coeff_atom_ener = 1
+            coeff_ener = 1
         # load energy_file
         else:
             coeff_ener, ener = self.load_data(set_name, energy_file, [nframes], False)