Update app.py

app.py

@@ -33,9 +33,10 @@ def classify_text(inp):
     with torch.no_grad():
         logits = model(**inputs).logits
     print(logits)
-    logits = torch.nn.Softmax(dim=0)(logits)
-    print(logits)
-    return {labels[i]: float(logits[i].item()) for i in range(len(labels))}
+    # logits = torch.nn.Softmax(dim=0)(logits)
+    # print(logits)
+    # return {labels[i]: float(logits[i].item()) for i in range(len(labels))}
+    return logits.argmax().item()
 
 gr.Interface(
     classify_text,