Update app.py

app.py

@@ -65,12 +65,19 @@ def predict(image):
 
     with torch.no_grad():  # Use no_grad context for inference to save memory and computations
         x = model(x)
-        probabilities = torch.nn.functional.softmax(x, dim=1)
-        class_id = probabilities.argmax(dim=1).item()
+        probabilities = torch.nn.functional.softmax(x, dim=1)[0]
+        #class_id = probabilities.argmax(dim=1).item()
+        cat_prob = probabilities[0]
+        dog_prob = probabilities[1]
 
-    classes = ['cat', 'dog']
-    return classes[class_id]
+    return {
+        'cat': cat_prob.item(),
+        'dog': dog_prob.item()
+    }
+
+    #classes = ['cat', 'dog']
+    #return classes[class_id]
 
 # Update Gradio interface
-demo = gr.Interface(fn=predict, inputs="image", outputs="text")
+demo = gr.Interface(fn=predict, inputs="image", outputs="label")
 demo.launch()