Update app.py

app.py

@@ -8,6 +8,14 @@ from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
 
 # Load Models
 def load_model(model_path, backbone_name, num_classes):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if backbone_name == "resnet50":
+        model = torch.load(model_path, map_location=device)
+    elif backbone_name == "mobilenet":
+        model = torch.load(model_path, map_location=device)
+    model.to(device)
+    model.eval()
+    return model
     if backbone_name == "resnet50":
         model = torch.load(model_path)
     elif backbone_name == "mobilenet":
@@ -79,4 +87,4 @@ video_interface = gr.Interface(
     title="Video Inference"
 )
 
-gr.TabbedInterface([image_interface, video_interface], ["Image", "Video"]).launch()
+gr.TabbedInterface([image_interface, video_interface], ["Image", "Video"]).launch()