diff --git a/local_integration_msms.py b/local_integration_msms.py
index 0015141dd22bf5c8f62930ca9a2aea60aa34307b..7cc384bc39cb6235c155158c80e3bded7679c839 100644
--- a/local_integration_msms.py
+++ b/local_integration_msms.py
@@ -1,20 +1,21 @@
 import pyopenms as oms
 import numpy as np
+import matplotlib.pyplot as plt
 
-def compute_chromatograms(rt,charge,intensity,start_c,end_c):
+def compute_chromatograms(rt, mz, intensity, start_c, end_c):
     value=[]
 
     for k in range(len(rt)):
-        c = np.array(charge[k])
+        c = np.array(mz[k])
         i = np.array(intensity[k])
-        value.append(np.sum(np.where(end_c > c > start_c, i, 0)))
+        value.append(np.sum(np.where((end_c > c) & (c > start_c), i, 0)))
 
     return value
 
 
 if __name__ == "__main__":
     e = oms.MSExperiment()
-    oms.MzMLFile().load("data/STAPH140.mzML", e)
+    oms.MzMLFile().load("data/Staph140.mzML", e)
     e.updateRanges()
     rt = []
     charge = []
@@ -24,4 +25,16 @@ if __name__ == "__main__":
             rt.append(s.getRT())
             charge.append(s.get_peaks()[0])
             intensity.append(s.get_peaks()[1])
-    val = compute_chromatograms(rt, charge, intensity, 400. ,400.5)
\ No newline at end of file
+    mz_range = np.linspace(350,1250,4000)
+    for i in range(len(mz_range)-1):
+        print(mz_range[i],'/1250')
+        val = compute_chromatograms(rt, charge, intensity, mz_range[i] ,mz_range[i+1])
+        fig, ax = plt.subplots()
+        ax.plot(val)
+        ax.set_xlabel('Retention time')
+        ax.set_ylabel('Intensity')
+        ax.set_title('mz : {} to {}'.format(mz_range[i] ,mz_range[i+1]))
+        plt.savefig('fig/rt_local/{}_to_{}.png'.format(mz_range[i] ,mz_range[i+1]))
+        plt.clf()
+df = e.get_df()
+#358.1 358.32
\ No newline at end of file
diff --git a/main_custom.py b/main_custom.py
index 1118c9aad5b7832c9bc62d7180de04ece7454e16..cca8c9ce9084ed12cfb64c0f6378451458a581ba 100644
--- a/main_custom.py
+++ b/main_custom.py
@@ -207,7 +207,7 @@ def main(args):
 
     print('\nData loaded')
 
-    model = Model_Common_Transformer_TAPE(encoder_ff=args.encoder_ff, decoder_rt_ff=args.decoder_rt_ff,
+    model = Model_Common_Transformer(encoder_ff=args.encoder_ff, decoder_rt_ff=args.decoder_rt_ff,
                                      decoder_int_ff=args.decoder_int_ff
                                      , n_head=args.n_head, encoder_num_layer=args.encoder_num_layer,
                                      decoder_int_num_layer=args.decoder_int_num_layer,