diff --git a/vgsales-mlflow.py b/vgsales-mlflow.py
index 3799c52..219de18 100644
--- a/vgsales-mlflow.py
+++ b/vgsales-mlflow.py
@@ -29,7 +29,7 @@ def my_main(epochs, batch_size):
     X_train, X_test, y_train, y_test = train_test_split(X,Y , test_size=0.2,train_size=0.8, random_state=21)
 
     model = Sequential()
-    model.add(Dense(9, input_dim = X_train, kernel_initializer='normal', activation='relu'))
+    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)