diff --git a/Zadanie_06_training.py b/Zadanie_06_training.py
index 7c7d342..98c9709 100644
--- a/Zadanie_06_training.py
+++ b/Zadanie_06_training.py
@@ -4,6 +4,7 @@ from sklearn.metrics import accuracy_score, classification_report
 import pandas as pd
 from sklearn.model_selection import train_test_split
 import numpy as np
+import sys
 from sklearn.preprocessing import StandardScaler, LabelEncoder
 from tensorflow.keras.optimizers import Adam
 
@@ -45,6 +46,6 @@ rms = Adam(lr=0.0003)
 
 NeuralModel.compile(optimizer=rms, loss='binary_crossentropy', metrics=['accuracy'])
 
-NeuralModel.fit(scaled_x, yenc, epochs=15) #verbose = 1
+NeuralModel.fit(scaled_x, yenc, batch_size=int(sys.argv[1]), epochs = int(sys.argv[1])) #verbose = 1
 
 NeuralModel.save('wine_model.h5')