Update app.py

app.py

@@ -11,8 +11,8 @@ model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
 
 def predict(image):
     inputs = feature_extractor(images=image, return_tensors="pt").to(device)
-    outputs = model(**inputs).to(device)
-    logits = outputs.logits.to(device)
+    outputs = model(**inputs)
+    logits = outputs.logits
     predicted_class_prob = F.softmax(logits, dim=-1).detach().numpy().max()
     predicted_class_idx = logits.argmax(-1).item()
     label = model.config.id2label[predicted_class_idx].split(",")[0]