chore: use cuda

app.py

@@ -16,6 +16,7 @@ MODEL_NAME = MODEL_NAMES[0]
 
 model = AutoModelForImageClassification.from_pretrained(
     MODEL_NAME,
+    device="cuda:0" if torch.cuda.is_available() else "cpu",
 )
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
 
@@ -90,7 +91,6 @@ def animagine_prompt(rating: list[str], character: list[str], general: list[str]
 
 
 @spaces.GPU
-@torch.no_grad()
 def predict_tags(
     image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8
 ):