Update app.py

app.py

@@ -40,8 +40,11 @@ if len(uploaded_files) == 0:
 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])
   predictions = cnn.apply({"params": params}, input)
+  st.write("Predictions shape: " + str(predictions.shape))
   for (image, prediction) in zip(uploaded_files, predictions):
     st.image(Image.open(image))
+    st.write("Prediction shape: " + str(prediction.shape))
+    st.write("Prediction argmax: " + str(prediction.argmax(axis=0)))
     st.write("Model Prediction: " + prediction.argmax(axis=0)[0])