diff --git a/JenkinsfileTraining b/JenkinsfileTraining
index 803474d..941a562 100644
--- a/JenkinsfileTraining
+++ b/JenkinsfileTraining
@@ -9,6 +9,11 @@ pipeline {
             description: 'Which build to use for copying artifacts',
             name: 'BUILD_SELECTOR'
         )
+        string(
+            defaultValue: '5',
+            description: 'Epochs',
+            name: 'EPOCHS'
+        )
     }
 
     triggers {
@@ -32,7 +37,7 @@ pipeline {
         stage('Script') {
             steps {
                 sh 'chmod 777 ./create_model.py'
-                sh "python3 ./create_model.py"
+                sh "python3 ./create_model.py ${params.EPOCHS}"
             }
         }
         stage('CreateArtifacts') {
diff --git a/create_model.py b/create_model.py
index a10823c..97b27d3 100644
--- a/create_model.py
+++ b/create_model.py
@@ -1,4 +1,5 @@
 import pandas as pd
+import sys
 from keras.models import Sequential
 from keras.layers import Dense
 from keras.optimizers import Adam
@@ -6,6 +7,8 @@ from keras import regularizers
 
 from helper import prepare_tensors
 
+epochs = int(sys.argv[1])
+
 hp_train = pd.read_csv('hp_train.csv')
 hp_dev = pd.read_csv('hp_dev.csv')
 
@@ -22,6 +25,6 @@ model.add(Dense(1, activation='linear'))
 adam = Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-7)
 model.compile(optimizer=adam, loss='mean_squared_error')
 
-model.fit(X_train, Y_train, epochs=20, batch_size=32, validation_data=(X_dev, Y_dev))
+model.fit(X_train, Y_train, epochs=epochs, batch_size=32, validation_data=(X_dev, Y_dev))
 
 model.save('hp_model.h5')