Update app.py

app.py

@@ -1,16 +1,18 @@
 from transformers import ViTFeatureExtractor, ViTForImageClassification
 from PIL import Image
+import torch
 import torch.nn.functional as F
 import time
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
-model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
+feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224').to(device)
+model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').to(device)
 
 def predict(image):
     inputs = feature_extractor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits
+    outputs = model(**inputs).to(device)
+    logits = outputs.logits.to(device)
     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]