fix: predict_fn return bug

app.py

@@ -19,7 +19,7 @@ transform = transforms.Compose([
         ])
 
 
-def predict_fn(image, raw_output=False):
+def predict_fn(image):
         """
         This function will predict the class of an image
         :param image_path: The path of the image
@@ -37,6 +37,7 @@ def predict_fn(image, raw_output=False):
         probabilities = torch.exp(output)
         top_p, top_class = probabilities.topk(1, dim=1)
 
-        return top_class.numpy()[0][0]
+        # return dictionary whose keys are labels and values are confidences
+        return {str(top_class.item()), str(top_p.item())}
 
 gr.Interface(predict_fn, gr.inputs.Image(type="pil", label="Input Image"), outputs="label").launch()