diff --git a/Zadanie_06_training.py b/Zadanie_06_training.py
index b6e51c8..6b908a8 100644
--- a/Zadanie_06_training.py
+++ b/Zadanie_06_training.py
@@ -68,7 +68,7 @@ def prepare_model(epoch_param, batch_param, _run):
     NeuralModel.save('wine_model.h5')
 
 @ex.main
-def my_main(train_size_param, test_size_param):
+def my_main(epoch_param, batch_param):
     print(prepare_model())