Update app.py

app.py

@@ -40,12 +40,9 @@ 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])
+    st.write("Model Prediction: " + prediction.argmax())
 
 
 def gridify(kernel, grid, kernel_size, scaling=5, padding=1):