Spaces:

p1atdev
/

wd-tagger-transformers

Running

p1atdev commited on Mar 21, 2024

Commit

4d398f3

1 Parent(s): 2068781

chore: forced to use cuda

Files changed (1) hide show

app.py CHANGED Viewed

@@ -16,7 +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)

 model = AutoModelForImageClassification.from_pretrained(
     MODEL_NAME,
+    device="cuda",
 )
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)