diff --git a/main_custom.py b/main_custom.py
index fec2d98427eae0a414849090dc94a461eb614c6b..a0f6489bd8271c46cd7d44e0243f71c181c9ade4 100644
--- a/main_custom.py
+++ b/main_custom.py
@@ -200,8 +200,8 @@ def run(epochs, eval_inter, save_inter, model, data_train, data_val, data_test,
 
     else :
         for e in range(1, epochs + 1):
-            # train(model, data_train, e, optimizer, criterion_rt, criterion_intensity, metric_rt, metric_intensity, forward,
-            #       wandb=wandb)
+            train(model, data_train, e, optimizer, criterion_rt, criterion_intensity, metric_rt, metric_intensity, forward,
+                  wandb=wandb)
             if e % eval_inter == 0:
                 eval(model, data_val, e, criterion_rt, criterion_intensity, metric_rt, metric_intensity, forward,
                      wandb=wandb)
diff --git a/prosit_data_merge.py b/prosit_data_merge.py
index 444200cb5ee06ac4bbbd467ddcbe14788b3e1fcc..1211691f84bfc89365ff2e314231465457e65752 100644
--- a/prosit_data_merge.py
+++ b/prosit_data_merge.py
@@ -148,7 +148,11 @@ def alphabetical_to_numerical(seq):
 
 df = pd.read_pickle('database/data_prosit_merged_holdout.pkl')
 
-df = df.head(100)
+print(len(df))
+print(df.head())
+df = df[df['Retention time']!=[0,0,0,0,0,0]]
+print(len(df))
+print(df.head())
 
-df.to_csv('database/data_head.csv')
+df.to_pickle('database/data_prosit_merged_holdout_2.pkl')