Solving Bug

app.py

@@ -21,7 +21,7 @@ if st.button("Predict"):
             outputs = model(**inputs)
             probs = F.softmax(outputs.logits, dim=-1)
             top = torch.argmax(probs, dim=1).item()
-            label = model.config.id2label[str(top)]
+            label = model.config.id2label[top]
             score = probs[0][top].item()
             
             st.success(f"**Prediction**: {label} ({score:.4f})")