diff --git a/model.py b/model.py
index d8bad31..0a34523 100644
--- a/model.py
+++ b/model.py
@@ -8,7 +8,9 @@ from tensorflow.keras.models import Sequential
 from tensorflow.keras.layers import Dense
 import tensorflow as tf
 
-data = pd.read_csv('./data/train.csv')
+# data = pd.read_csv('./data/train.csv')
+data = pd.read_csv('./openpowerlifting.csv')
+print(data)
 
 data = data[['Sex', 'Age', 'BodyweightKg', 'TotalKg']].dropna()
 
@@ -21,7 +23,7 @@ preprocessor = ColumnTransformer(
     transformers=[
         ('num', StandardScaler(), ['Age', 'BodyweightKg']),
         ('cat', OneHotEncoder(), ['Sex'])
-    ]
+    ],
 )
 
 pipeline = Pipeline(steps=[
@@ -35,8 +37,9 @@ pipeline = Pipeline(steps=[
 
 pipeline['model'].compile(optimizer='adam', loss='mse', metrics=['mae'])
 
-X_train_excluded = X_train.iloc[:, 1:]
+X_train_excluded = X_train.iloc[1:]
+y_train_excluded = y_train.iloc[1:]
 
-pipeline.fit(X_train_excluded, y_train, model__epochs=10, model__validation_split=0.1)
+pipeline.fit(X_train_excluded, y_train_excluded, model__epochs=int(sys.argv[1]), model__validation_split=0.1)
 
 pipeline['model'].save('powerlifting_model.h5')