diff --git a/Zadanie_06_training.py b/Zadanie_06_training.py
index 98c9709..96f6c3d 100644
--- a/Zadanie_06_training.py
+++ b/Zadanie_06_training.py
@@ -46,6 +46,6 @@ rms = Adam(lr=0.0003)
 
 NeuralModel.compile(optimizer=rms, loss='binary_crossentropy', metrics=['accuracy'])
 
-NeuralModel.fit(scaled_x, yenc, batch_size=int(sys.argv[1]), epochs = int(sys.argv[1])) #verbose = 1
+NeuralModel.fit(scaled_x, yenc, batch_size=int(sys.argv[1]), epochs = int(sys.argv[2])) #verbose = 1
 
 NeuralModel.save('wine_model.h5')