go back to previous probability computation

app.py

@@ -34,11 +34,12 @@ def predict_fn(image):
         with torch.no_grad():
             output = model(image)
         
-        print(output)
+        probabilities = torch.exp(output)
+        print(probabilities)
         # focused (0) round probability
-        focused_prob = round(output[0][0].item(), 4)
+        focused_prob = round(probabilities[0][0].item(), 2)
         # distracted (1) round probability
-        distracted_prob = round(output[0][1].item(), 4)
+        distracted_prob = round(probabilities[0][1].item(), 2)
 
         # return dictionary whose keys are labels and values are confidences
         return {'focused': focused_prob, 'distracted': distracted_prob}