Update modelutil.py (#1)

- Update modelutil.py (c9c13224e4b1d7f7de5f2133067515ea62d0020b)


Co-authored-by: papa wade <papasega@users.noreply.huggingface.co>

modelutil.py

@@ -2,7 +2,9 @@ import tensorflow as tf
 
 def create_model():
 
-    LAYERS = [tf.keras.layers.Flatten(input_shape=[28,28], name="inputlayer"),
+    LAYERS = [tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1), name="convlayer1"),
+              tf.keras.layers.Conv2D(64, (3, 3), activation='relu', name="convlayer2"),
+              tf.keras.layers.Flatten(input_shape=[28,28], name="inputlayer"),
               tf.keras.layers.Dense(300, activation='relu', name="hiddenlayer1"),
               tf.keras.layers.Dense(100, activation='relu', name="hiddenlayer2"),
               tf.keras.layers.Dense(10, activation='softmax', name="outputlayer")]