diff --git a/toolkit.py b/toolkit.py
index ec09631b541986c761c610ad3eebb02184ca4893..70aa273bfc095469bd66aeb9d1f4d78e4b18d4f7 100755
--- a/toolkit.py
+++ b/toolkit.py
@@ -560,8 +560,8 @@ class KerasROOTClassification(object):
             # join scores and is_train array
             scores = self.evaluate(x_eval).reshape(-1)
             friend_df = pd.DataFrame(np.array(scores, dtype=[(score_name, np.float64)]))
-            friend_df["is_train"] = is_train
-            friend_tree = friend_df.to_records()[[score_name, "is_train"]]
+            friend_df[score_name+"_is_train"] = is_train
+            friend_tree = friend_df.to_records()[[score_name, score_name+"_is_train"]]
             if start == 0:
                 mode = "recreate"
             else:
@@ -742,7 +742,6 @@ if __name__ == "__main__":
                                 identifiers = ["DatasetNumber", "EventNumber"],
                                 step_bkg = 100)
 
-    c.load()
     c.train(epochs=20)
     c.plot_ROC()
     c.plot_loss()