Update app.py

app.py

@@ -39,10 +39,10 @@ if len(uploaded_files) == 0:
     st.write("Please upload an image!")
 else:
   input = jnp.array([jnp.array(Image.open(uploaded_file).convert("HSV").split()[2].resize((28, 28))).reshape(1, 28, 28, 1) / 255. for uploaded_file in uploaded_files])
-  prediction = cnn.apply({"params": params}, input)
-  for (index, image) in enumerate(uploaded_files):
+  predictions = cnn.apply({"params": params}, input)
+  for (image, prediction) in zip(uploaded_files, predictions):
     st.image(Image.open(image))
-    st.write("Model Prediction: " + prediction[index].argmax(axis=0)[0])
+    st.write("Model Prediction: " + prediction.argmax(axis=0)[0])
 
 
 def gridify(kernel, grid, kernel_size, scaling=5, padding=1):