Update app.py

app.py

@@ -6,7 +6,11 @@ from gradio import inputs, outputs
 from PIL import Image
 from torchvision import transforms
 
-model = torch.hub.load("RF5/danbooru-pretrained", "resnet50")
+model = torch.hub.load(
+    "RF5/danbooru-pretrained",
+    "resnet50",
+    map_location=torch.device("cpu"),
+)
 model.eval()
 
 with open("./tags.json", "rt", encoding="utf-8") as f:
@@ -28,10 +32,6 @@ def main(input_image: Image.Image, threshold: float):
         0
     )  # create a mini-batch as expected by the model
 
-    if torch.cuda.is_available():
-        input_batch = input_batch.to("cuda")
-        model.to("cuda")
-
     with torch.no_grad():
         output, *_ = model(input_batch)