feat: add cuda support

app.py

@@ -12,7 +12,7 @@ import torchvision.transforms as T
 
 logging.getLogger('PIL').setLevel(logging.CRITICAL)
 
-device = torch.device("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 def load_model(baseline: bool = False):
@@ -35,7 +35,9 @@ def frcnn_motsynth(image):
     model = load_model()
     transformEval = presets.DetectionPresetEval()
     image_tensor = transformEval(image, None)[0]
+    image_tensor = image_tensor.to(device)
     prediction = model([image_tensor])[0]
+    prediction = prediction.to(device)
     image_w_bbox = add_bbox(image_tensor, prediction, 0.80)
     torchvision.io.write_png(image_w_bbox, "custom_out.png")
     return "custom_out.png"
@@ -45,7 +47,9 @@ def frcnn_coco(image):
     model = load_model(baseline=True)
     transformEval = presets.DetectionPresetEval()
     image_tensor = transformEval(image, None)[0]
+    image_tensor = image_tensor.to(device)
     prediction = model([image_tensor])[0]
+    prediction = prediction.to(device)
     image_w_bbox = add_bbox(image_tensor, prediction, 0.80)
     torchvision.io.write_png(image_w_bbox, "baseline_out.png")
     return "baseline_out.png"