Update app.py

app.py

@@ -2,10 +2,21 @@ import json
 
 import gradio
 import torch
+import torch.hub
 from gradio import inputs, outputs
 from PIL import Image
 from torchvision import transforms
 
+real_load = torch.hub.load_state_dict_from_url
+
+
+def load_state_dict_from_url(*args, **kwargs):
+    kwargs["map_location"] = "cpu"
+    return real_load(*args, **kwargs)
+
+
+torch.hub.load_state_dict_from_url = load_state_dict_from_url
+
 model = torch.hub.load(
     "RF5/danbooru-pretrained",
     "resnet50",