Update app.py

app.py

@@ -79,8 +79,7 @@ def predict(img):
         output = model(img_input)
         pred_class = torch.argmax(output, dim=1).item()
 
-    return {class_ids[int2id[pred_class]]: 1.0}
-
+    return {class_ids[pred_class]: 1.0}
 # ----------------------------
 # Gradio UI
 # ----------------------------