Update app.py

app.py

@@ -1,12 +1,9 @@
-from transformers import AutoModel
+from transformers import AutoModelForImageClassification, ViTImageProcessor
 import gradio as gr
 from PIL import Image
 import torch
-from transformers import AutoImageProcessor, AutoModelForImageClassification
-from transformers import ViTImageProcessor, ViTConfig, ViTForImageClassification
-from transformers import TrainingArguments, Trainer
 
-model = AutoModel.from_pretrained("KFrimps/oxford-pets-vit-from-scratch")
+model = AutoModelForImageClassification.from_pretrained("KFrimps/oxford-pets-vit-from-scratch")
 processor = ViTImageProcessor.from_pretrained("KFrimps/oxford-pets-vit-from-scratch")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"