sync with remote

app.py

@@ -123,10 +123,11 @@ tf.keras.backend.clear_session()
 
 # set up the model architecture
 model = tf.keras.models.Sequential([
-    Flatten(input_shape=[256, 256, 1]),
-    Dense(64, activation='relu'),
-    Dense(256*256*2, activation='softmax'),
-    Reshape((256, 256, 2))
+    tf.keras.layers.Input(shape=(256, 256, 1)),  # Define input shape
+    tf.keras.layers.Flatten(),
+    tf.keras.layers.Dense(64, activation='relu'),
+    tf.keras.layers.Dense(256*256*2, activation='softmax'),
+    tf.keras.layers.Reshape((256, 256, 2))
 ])
 
 # specify how to train the model with algorithm, the loss function and metrics
@@ -140,4 +141,7 @@ model_summary = StringIO()
 model.summary(print_fn=lambda x: model_summary.write(x + '\n'))
 
 # Display the model summary in Streamlit
-st.markdown(model_summary.getvalue())
+st.markdown(model_summary.getvalue())
+
+# plot the model including the sizes of the model
+tf.keras.utils.plot_model(model, show_shapes=True)