diff --git a/train.py b/train.py
index e54ae26..fc96300 100644
--- a/train.py
+++ b/train.py
@@ -15,15 +15,15 @@ exint.observers.append(FileStorageObserver('my_runs'))
 def my_config():
     EPOCHS = int(os.environ['EPOCHS'])
     train_data_x = exint.open_resource('./X_train.csv', 'r')
-    games_all = train_data_x.copy()
-    games_predict = train_data_x.pop('User_Score')
+
     normalize = layers.Normalization()
     _run.info["epochs"] = EPOCHS
 
 
 @exint.main
-def my_main(EPOCHS, games_all, games_predict,normalize):
-    _run.info["prepare_message_ts"] = str(datetime.now())
+def my_main(EPOCHS,train_data_x, normalize):
+    games_all = train_data_x.copy()
+    games_predict = train_data_x.pop('User_Score')
     normalize.adapt(games_all)
     norm_games_model = tensorflow.keras.Sequential([
         normalize,