Update app.py

app.py

@@ -119,6 +119,8 @@ model = timm.create_model(
 ) # type: VisionTransformer
 
 safetensors.torch.load_model(model, "JTP_PILOT-e4-vit_so400m_patch14_siglip_384.safetensors")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 model.eval()
 
 with open("tagger_tags.json", "r") as file:
@@ -135,7 +137,7 @@ def run_classifier(image, threshold):
     global sorted_tag_score
     img = image.convert('RGB')
     tensor = transform(img).unsqueeze(0)
-
+    tensor = tensor.to(device)
     with torch.no_grad():
         logits = model(tensor)
         probabilities = torch.nn.functional.sigmoid(logits[0])