Update modelutil.py

modelutil.py

@@ -1,20 +1,24 @@
 import tensorflow as tf
 
-LAYERS = [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")]
+def create_model():
 
-model = tf.keras.models.Sequential(LAYERS)
+    LAYERS = [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")]
+    
+    model = tf.keras.models.Sequential(LAYERS)
+    
+    model.load_weights('./checkpoint')
+    
+    
+    LOSS_FUNCTION = tf.keras.losses.SparseCategoricalCrossentropy() # HERE
+    OPTIMIZER = "ADAM"
+    METRICS = ["accuracy"]
+    model.compile(loss=LOSS_FUNCTION,
+    optimizer=OPTIMIZER,
+    metrics=METRICS)
 
-model.load_weights('./checkpoint')
-
-
-LOSS_FUNCTION = tf.keras.losses.SparseCategoricalCrossentropy() # HERE
-OPTIMIZER = "ADAM"
-METRICS = ["accuracy"]
-model.compile(loss=LOSS_FUNCTION,
-optimizer=OPTIMIZER,
-metrics=METRICS)
+    return model