diff --git a/train.py b/train.py
index 6d2ae90..67d14d7 100755
--- a/train.py
+++ b/train.py
@@ -26,9 +26,10 @@ def my_config():
     batch_param = int(sys.argv[1])
     epoch_param = int(sys.argv[2])
 
+@ex.capture
 def prepare_model(epoch_param, batch_param, _run):
     _run.info["prepare_model_ts"] = str(datetime.now())
-    # odczytanie danych z plików
+
     vgsales_train = pd.read_csv('train.csv')
     vgsales_test = pd.read_csv('test.csv')
     vgsales_dev = pd.read_csv('dev.csv')
@@ -37,33 +38,35 @@ def prepare_model(epoch_param, batch_param, _run):
     vgsales_test['Nintendo'] = vgsales_test['Publisher'].apply(lambda x: 1 if x=='Nintendo' else 0)
     vgsales_dev['Nintendo'] = vgsales_dev['Publisher'].apply(lambda x: 1 if x=='Nintendo' else 0)
 
-    # podzial na X i y
     X_train = vgsales_train.drop(['Rank','Name','Platform','Year','Genre','Publisher'],axis = 1)
     y_train = vgsales_train[['Nintendo']]
     X_test = vgsales_test.drop(['Rank','Name','Platform','Year','Genre','Publisher'],axis = 1)
     y_test = vgsales_test[['Nintendo']]
 
     print(X_train.shape[1])
-    # keras model
+
     model = Sequential()
     model.add(Dense(9, input_dim = X_train.shape[1], kernel_initializer='normal', activation='relu'))
     model.add(Dense(1,kernel_initializer='normal', activation='sigmoid'))
 
     early_stop = EarlyStopping(monitor="val_loss", mode="min", verbose=1, patience=10)
 
-    # kompilacja
     model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
 
-    # model fit
     epochs = int(sys.argv[1])
     batch_size = int(sys.argv[2])
 
-    # trenowanie modelu
     model.fit(X_train, y_train, epochs=epochs, batch_size=batch_size, validation_data=(X_test, y_test))
 
-    # zapisanie modelu
+    prediction = model.predict(X_test)
+
+    rmse = mean_squared_error(y_test, prediction)
+    _run.log_scalar("rmse", rmse)
+
     model.save('vgsales_model.h5')
 
+    return rmse
+
 @ex.main
 def my_main(epoch_param, batch_param):
     print(prepare_model())