Update app.py

app.py

@@ -10,7 +10,7 @@ feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16
 model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').to(device)
 
 def predict(image):
-    inputs = feature_extractor(images=image, return_tensors="pt")
+    inputs = feature_extractor(images=image, return_tensors="pt").to(device)
     outputs = model(**inputs).to(device)
     logits = outputs.logits.to(device)
     predicted_class_prob = F.softmax(logits, dim=-1).detach().numpy().max()