Update app.py

app.py

@@ -31,7 +31,8 @@ predictTransform = transforms.Compose([
 
 def predict(img):
     img = predictTransform(img).unsqueeze(0).to(DEVICE)
-    with torch.inference_mode():
+    with torch.no_grad():
+        model.eval()
         prediction = torch.nn.functional.softmax(model(img)[0], dim=0)
         confidences = {labels[i]: float(prediction[i]) for i in range(len(labels))}
     return confidences