change in probs computation

app.py

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