diff --git a/Armory/TrainingBuilder.py b/Armory/TrainingBuilder.py
index 8ccb613..e360b5d 100644
--- a/Armory/TrainingBuilder.py
+++ b/Armory/TrainingBuilder.py
@@ -32,7 +32,7 @@ from tensorflow.keras import Sequential
 from tensorflow.keras.layers import Dense
 from tensorflow.keras.callbacks import EarlyStopping
 
-ROOT_FOLDER = Path("/Users/jamesszalkie/ANASEN_analysis/Armory/Training_Data_Ne/")
+ROOT_FOLDER = Path("/Users/jamesszalkie/ANASEN_analysis/Armory/Training_Data/")
 
 TREE_NAME = "tree1"
 
@@ -44,6 +44,10 @@ INPUT_BRANCHES = [
     "MTarget",
     "MLight",
     "MHeavy",
+    "ZBeam",
+    "ZHeavy",
+    "ZLight",
+    "ZTarget"
 ]
 
 OUTPUT_BRANCHES = [