Update app.py

app.py

@@ -136,8 +136,8 @@ prediction = np.argmax(prediction, axis=1)
 st.write(prediction)
 st.write(targets)
 
-st.write("after argmax axis 0")
-targets = np.argmax(targets, axis=0)
-prediction = np.argmax(prediction, axis=0)
+st.write("after argmax axis 2")
+targets = np.argmax(targets, axis=2)
+prediction = np.argmax(prediction, axis=2)
 st.write(prediction)
 st.write(targets)