Update app.py

app.py

@@ -43,10 +43,10 @@ if len(uploaded_files) == 0:
   st.write("Please upload an image!")
 else:
   input = jnp.array([tf.cast(tf.image.resize(tf.convert_to_tensor(Image.open(uploaded_file)), [50, 50]), tf.float32) / 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))
-    [cat_prob, dog_prob] = jax.nn.softmax(prediction[index])
+    [cat_prob, dog_prob] = jax.nn.softmax(prediction)
     if cat_prob > dog_prob:
       st.write(f"Model Prediction - Cat ({100*cat_prob:.2f}%), Dog ({100*dog_prob:.2f}%)")
     else: