Update app.py

app.py

@@ -7,10 +7,13 @@ from transformers import AutoModelForImageClassification, AutoConfig
 import gradio as gr
 import spaces
 
+
+
 model_id = "thelabel/240903-image-tagging"
 config = AutoConfig.from_pretrained(model_id)
 model = AutoModelForImageClassification.from_pretrained(model_id)
-model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
 # Standard ViT image transforms
 image_transform = transforms.Compose([
@@ -33,7 +36,7 @@ def predict_tags(image_url, threshold=0.5):
     if image is None:
         return [], "Could not load image from the provided URL."
 
-    image_tensor = image_transform(image).unsqueeze(0)
+    image_tensor = image_transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
         logits = model(image_tensor).logits
         probs = torch.sigmoid(logits).squeeze()