Update main.py

main.py

@@ -43,6 +43,7 @@ def preprocess():
 def train():
     input_shape = (28, 28, 1)
     x_train, y_train, x_test, y_test = preprocess()
+    num_classes = 10
 
     model = keras.Sequential(
         [