Update app.py

app.py

@@ -1,10 +1,16 @@
 from transformers import AutoModel
+import gradio as gr
+from PIL import Image
+import torch
 
 model = AutoModel.from_pretrained("KFrimps/oxford-pets-vit-from-scratch")
+processor = ViTImageProcessor.from_pretrained("KFrimps/oxford-pets-vit-from-scratch")
 
-import gradio as gr
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+id2label = model.config.id2label
 
-from PIL import Image
 
 def predict(image):
   """Predicts the class of the input image using the fine-tuned student model."""