Update handler.py

handler.py

@@ -21,6 +21,6 @@ class EndpointHandler:
             inputs = [inputs]
 
         # run normal prediction
-        scores = self.model.predict(inputs)[0]
+        scores = self.model.predict_proba(inputs)[0]
 
         return [{"label": self.id2label[i], "score": score.item()} for i, score in enumerate(scores)]