Update app.py

app.py

@@ -53,24 +53,24 @@ if st.button("Train"):
     with col4:
         st.write("Conv2D Layer-2")
         sub1 = Model(inputs=model.inputs[0],outputs=model.layers[2].output)
-        fig,ax2 = plt.subplots(6,1,figsize=(8,6))
-        for i in range(6):
+        fig,ax2 = plt.subplots(16,1,figsize=(8,6))
+        for i in range(16):
             ax2[i].imshow(sub1.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
             ax2[i].axis("off")
         st.pyplot(fig)
     with col5:
         st.write("Average Pooling Layer-2")
         max2 = Model(inputs=model.inputs[0],outputs=model.layers[3].output)
-        fig,ax3 = plt.subplots(6,1,figsize=(8,6))
-        for i in range(6):
+        fig,ax3 = plt.subplots(16,1,figsize=(8,6))
+        for i in range(16):
             ax3[i].imshow(max2.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
             ax3[i].axis("off")
         st.pyplot(fig)
     with col6:
-        st.write("Average Pooling Layer-2")
+        st.write("Conv2D Layer-3")
         sub3 = Model(inputs=model.inputs[0],outputs=model.layers[4].output)
-        fig,ax4 = plt.subplots(6,1,figsize=(8,6))
-        for i in range(6):
+        fig,ax4 = plt.subplots(120,1,figsize=(8,6))
+        for i in range(120):
             ax4[i].imshow(sub3.predict(x_train[0:1,:].reshape(1,28,28,1))[0,:,:,i],cmap="gray")
             ax4[i].axis("off")
         st.pyplot(fig)