diff --git a/Zadanie_06_training.py b/Zadanie_06_training.py
index c2e7c41..1997aa1 100644
--- a/Zadanie_06_training.py
+++ b/Zadanie_06_training.py
@@ -31,6 +31,6 @@ def regression_model():
     return model
 
 model = regression_model()
-model.fit(x_train, y_train, batch_size=int(sys.argv[1]), epochs = sys.argv[2]) #verbose = 1
+model.fit(x_train, y_train, batch_size=int(sys.argv[1]), epochs = int(sys.argv[1])) #verbose = 1
 
 model.save('wine_model')
\ No newline at end of file