[add] device

run.py

@@ -62,7 +62,7 @@ def inference(image_pil):
     backbone = resnet_fpn_backbone('resnet18', False)
     model = FasterRCNN(backbone, num_classes)
     model.eval()
-    state_dict = torch.load('model/model/densenet-model-9-mAp--1.0.pth')
+    state_dict = torch.load('model/model/densenet-model-9-mAp--1.0.pth',map_location=device)
     model.load_state_dict(state_dict["model"])
     _transform = T.Compose([T.ToTensor()])
     image = image_pil.convert("RGB")