Update app.py

app.py

@@ -32,6 +32,7 @@ if st.button("Train"):
                 cmap="gray"
             )
             axs[i].axis("off")
+        st.pyplot(fig)
     with col2:
         st.write("Conv2D Layer")
         sub = Model(inputs=model.inputs[0],outputs=model.layers[0].output)
@@ -39,9 +40,12 @@ if st.button("Train"):
         for i in range(6):
             ax[i].imshow(sub.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
             ax[i].axis("off")
+        st.pyplot(fig)
     with col3:
         st.write("Average Pooling Layer")
         max1 = Model(inputs=model.inputs[0],outputs=model.layers[1].output)
         fig,ax1 = plt.subplots(6,1,figsize=(8,6))
         for i in range(6):
             ax1[i].imshow(max1.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
+            ax1[i].axis("off")
+        st.pyplot(fig)