Update app.py

app.py

@@ -13,7 +13,7 @@ def predict(image):
     inputs = feature_extractor(images=image, return_tensors="pt").to(device)
     outputs = model(**inputs)
     logits = outputs.logits
-    predicted_class_prob = F.softmax(logits, dim=-1).detach().numpy().max()
+    predicted_class_prob = F.softmax(logits, dim=-1).detach().cpu().numpy().max()
     predicted_class_idx = logits.argmax(-1).item()
     label = model.config.id2label[predicted_class_idx].split(",")[0]
     time.sleep(2)