diff --git a/data_viz.py b/data_viz.py
index deabf69e28821d5d94066040ad7fb58b1d1d2840..312e24c1c309dadf498fde87d472205c4ec1f0b9 100644
--- a/data_viz.py
+++ b/data_viz.py
@@ -275,12 +275,12 @@ def add_length(dataframe):
     dataframe['length']=dataframe['seq'].map(fonc)
 
 
-df = pd.read_csv('output/out_common_ISA_augmented_10_eval.csv')
+df = pd.read_csv('output/out_ISA_augmented_2_align_eval.csv')
 add_length(df)
 df['abs_error'] =  np.abs(df['rt pred']-df['true rt'])
-histo_abs_error(df, display=False, save=True, path='fig/custom model res/histo_SA_augmented_10_eval.png')
-scatter_rt(df, display=False, save=True, path='fig/custom model res/RT_pred_SA_augmented_10_eval.png', color=True)
-histo_length_by_error(df, bins=10, display=False, save=True, path='fig/custom model res/histo_length_SA_augmented_10_eval.png')
+histo_abs_error(df, display=False, save=True, path='fig/custom model res/histo_ISA_augmented_2_align_eval.png')
+scatter_rt(df, display=False, save=True, path='fig/custom model res/RT_pred_ISA_augmented_2_align_eval.png', color=True)
+histo_length_by_error(df, bins=10, display=False, save=True, path='fig/custom model res/histo_length_ISA_augmented_2_align_eval.png')
 #
 # df = pd.read_csv('output/out_common_ISA_augmented_3_eval.csv')
 # add_length(df)
diff --git a/main_ray_tune.py b/main_ray_tune.py
index 85d17a7fc232b44743f6012e2af92990941a81e7..d939a1b85b99ec23e4f231470c349925bd6629b9 100644
--- a/main_ray_tune.py
+++ b/main_ray_tune.py
@@ -277,17 +277,17 @@ def main(args, gpus_per_trial=1):
     #     "batch_size": tune.choice([1024]),
     # }
     config = {
-        "encoder_num_layer": tune.choice([2, 4, 8]),
-        "decoder_rt_num_layer": tune.choice([2, 4, 8]),
-        "decoder_int_num_layer": tune.choice([1]),
+        "encoder_num_layer": tune.choice([2]),
+        "decoder_rt_num_layer": tune.choice([1]),
+        "decoder_int_num_layer": tune.choice([1,2,4]),
         "embedding_dim": tune.choice([16, 64]),
-        "encoder_ff": tune.choice([512, 1024, 2048]),
-        "decoder_rt_ff": tune.choice([512, 1024, 2048]),
-        "decoder_int_ff": tune.choice([512]),
+        "encoder_ff": tune.choice([2048]),
+        "decoder_rt_ff": tune.choice([2048]),
+        "decoder_int_ff": tune.choice([512, 1024, 2048]),
         "n_head": tune.choice([1, 2, 4, 8, 16]),
         "drop_rate": tune.choice([0.25]),
         "lr": tune.loguniform(1e-4, 1e-2),
-        "batch_size": tune.choice([4096]),
+        "batch_size": tune.choice([1024,2048,4096]),
         "activation": tune.choice(['relu', 'gelu']),
         "norm_first": tune.choice([True, False]),
     }