diff --git a/train.py b/train.py
index 67d14d7..d15a2bb 100755
--- a/train.py
+++ b/train.py
@@ -43,10 +43,8 @@ def prepare_model(epoch_param, batch_param, _run):
     X_test = vgsales_test.drop(['Rank','Name','Platform','Year','Genre','Publisher'],axis = 1)
     y_test = vgsales_test[['Nintendo']]
 
-    print(X_train.shape[1])
-
     model = Sequential()
-    model.add(Dense(9, input_dim = X_train.shape[1], kernel_initializer='normal', activation='relu'))
+    model.add(Dense(9, input_dim = X_train, 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)