Update app.py

app.py

@@ -4,12 +4,13 @@ import torch
 import gradio as gr
 from transformers import pipeline
 
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Loading in Model
 model_name = "imjeffhi/pokemon_classifier"
-model = ViTForImageClassification.from_pretrained(model_name).to(device)
 feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
-model.eval()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = ViTForImageClassification.from_pretrained(model_name).to(device)
 model.to(device)