diff --git a/architectures/utils/toolbox.py b/architectures/utils/toolbox.py
index 601eb14..67fa32f 100644
--- a/architectures/utils/toolbox.py
+++ b/architectures/utils/toolbox.py
@@ -4,6 +4,7 @@
 import tensorflow as tf
 import pickle
 import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
 
 rnd_seed=1
 
@@ -511,8 +512,8 @@
         X_batch_A=denormalise(train_true, min_true, max_true, norm_space=True)
         X_batch_B=denormalise(train_reco, min_reco, max_reco)
 
-    X_batch_A = X_batch_A[j]
-    X_batch_B = X_batch_B[j]
+    X_batch_A = train_true[j]
+    X_batch_B = train_reco[j]
 
     n_H_B, n_W_B, n_C = X_batch_B.shape
 
@@ -555,27 +556,50 @@
         X_A=denormalise(X_A, min_true, max_true)
         X_B=denormalise(X_B, min_reco, max_reco)
         sample_nn=denormalise(sample_nn, min_reco, max_reco)
-
+    
     plt.subplot(1,3,1)
-    plt.gca().set_title('True ET {0:.6g}'.format(X_A.sum()))
-    plt.imshow(X_A.reshape(n_H_A,n_W_A))
+    plt.gca().set_title('RealET from sim')
     plt.xlabel('X')
     plt.ylabel('Y')
-    plt.subplots_adjust(wspace=0.2,hspace=0.2)
+    ax = plt.gca()
+    im = ax.imshow(X_A.reshape(n_H_A,n_W_A)/1000)
+    # create an axes on the right side of ax. The width of cax will be 5%
+    # of ax and the padding between cax and ax will be fixed at 0.05 inch.
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+
+    plt.colorbar(im, cax=cax)
+    plt.subplots_adjust(wspace=0.4,hspace=0.2)
 
     plt.subplot(1,3,2)
-    plt.gca().set_title('MC Reco ET {0:.6g}'.format(X_B.sum()))
-    plt.imshow(X_B.reshape(n_H_B,n_W_B))
+    plt.gca().set_title('RecoET Geant 4')
+    #plt.imshow(X_B.reshape(n_H_B,n_W_B))
     plt.xlabel('X')
     plt.ylabel('Y')
-    plt.subplots_adjust(wspace=0.2,hspace=0.2)
+    ax = plt.gca()
+    im = ax.imshow(X_B.reshape(n_H_B,n_W_B)/1000)
+    # create an axes on the right side of ax. The width of cax will be 5%
+    # of ax and the padding between cax and ax will be fixed at 0.05 inch.
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+
+    plt.colorbar(im, cax=cax)
+    plt.subplots_adjust(wspace=0.4,hspace=0.2)
 
     plt.subplot(1,3,3)
-    plt.gca().set_title('NN Reco ET {0:.6g}'.format(sample_nn.sum()))
-    plt.imshow(sample_nn.reshape(n_H_B,n_W_B))
+    plt.gca().set_title('RecoET from NN')
+    #plt.imshow(sample_nn.reshape(n_H_B,n_W_B)/10)
     plt.xlabel('X')
     plt.ylabel('Y')
-    plt.subplots_adjust(wspace=0.2,hspace=0.2)
+    ax = plt.gca()
+    im = ax.imshow(sample_nn.reshape(n_H_B,n_W_B)/1000)
+    # create an axes on the right side of ax. The width of cax will be 5%
+    # of ax and the padding between cax and ax will be fixed at 0.05 inch.
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+
+    plt.colorbar(im, cax=cax)
+    plt.subplots_adjust(wspace=0.4,hspace=0.2)
     if is_training:
         plt.suptitle('At iter {0}'.format(total_iters))
     fig = plt.gcf()
@@ -583,9 +607,9 @@
 
     if save:
         if is_training:
-            plt.savefig(PATH+'/sample_at_iter_{0}.png'.format(total_iters),dpi=80)
+            plt.savefig(PATH+'/sample_at_iter_{0}.eps'.format(total_iters), format='eps', dpi=100)
         else:
-            plt.savefig(PATH+'/nn_reco_sample_{0}.png'.format(j),dpi=80)
+            plt.savefig(PATH+'/nn_reco_sample_{0}.eps'.format(j), format='eps',dpi=100)
     else:
         plt.show()
 
@@ -663,8 +687,8 @@
     value = np.zeros(shape=(2,2))
     pos =np.zeros(shape=(2,1))
     
-    for i in range(img.shape[0]-1):
-        for j in range(img.shape[1]-1):
+    for i in range(img.shape[0]-2):
+        for j in range(img.shape[1]-2):
             c_prime = img[i:i+2,j:j+2].sum()
             if c_prime > c:
                 c = c_prime
@@ -678,18 +702,27 @@
     return value, pos
 
 def get_triggered_events(true, reco_inner, reco_outer):
-    l = []
+    l_inner = []
+    l_outer = []
     for m in range(len(reco_inner)):
         value_inner, pos_inner = get_4_max_cells(reco_inner[m])
+        
+        if value_inner.sum()>3680: 
+                l_inner.append(m)
+
+    for m in range(len(reco_outer)):
         value_outer, pos_outer = get_4_max_cells(reco_outer[m])
         
-        if value_inner.sum()>3680 or value_outer.sum()>3680:
-                l.append(m)
-    triggered_true = np.array([true[l[i]].sum() for i in range(len(l))])
-    triggered_reco_inner = np.array([reco_inner[l[i]].sum() for i in range(len(l))])
-    triggered_reco_outer = np.array([reco_outer[l[i]].sum() for i in range(len(l))])
+        if value_outer.sum()>3680:
+                l_outer.append(m)
+
+    triggered_true_inner = np.array([true[l_inner[i]].sum() for i in range(len(l_inner))])
+    triggered_true_outer = np.array([true[l_outer[i]].sum() for i in range(len(l_outer))])
+
+    triggered_reco_inner = np.array([reco_inner[l_inner[i]].sum() for i in range(len(l_inner))])
+    triggered_reco_outer = np.array([reco_outer[l_outer[i]].sum() for i in range(len(l_outer))])
     
-    return l, triggered_true, triggered_reco_inner, triggered_reco_outer
+    return l_inner, l_outer, triggered_true_inner, triggered_true_outer, triggered_reco_inner, triggered_reco_outer
 
 # def crop_conditional(true, reco, dim):