a

app.py

@@ -53,7 +53,7 @@ def classify_image(img):
     img = tf.convert_to_tensor(img, dtype=tf.float32)
     img = tf.expand_dims(img, axis=0)
     # probs = model.predict(img).flatten()
-    probs = (model(img)).flatten()
+    probs = (model(img))[0]
     return {labels[i]: float(probs[i]) for i in range(len(labels))}
 
 image = gr.inputs.Image(shape=(256, 256))