Update main.py

main.py

@@ -34,14 +34,14 @@ def preprocess():
     x_test = np.expand_dims(x_test, -1)
 
     num_classes = 10
-    input_shape = (28, 28, 1)
-    
+
     y_train = keras.utils.to_categorical(y_train, num_classes)
     y_test = keras.utils.to_categorical(y_test, num_classes)
 
     return x_train, y_train, x_test, y_test
 
 def train():
+    input_shape = (28, 28, 1)
     x_train, y_train, x_test, y_test = preprocess()
 
     model = keras.Sequential(