Update my_model/object_detection.py

my_model/object_detection.py

@@ -66,7 +66,7 @@ class ObjectDetector:
         try:
             model_path = get_model_path('deformable-detr-detic')
             self.processor = AutoImageProcessor.from_pretrained(model_path)
-            self.model = AutoModelForObjectDetection.from_pretrained(model_path)
+            self.model = AutoModelForObjectDetection.from_pretrained(model_path, device_map = self.device)
         except Exception as e:
             print(f"Error loading Detic model: {e}")
             raise
@@ -84,9 +84,9 @@ class ObjectDetector:
         try:
             model_path = get_model_path ('yolov5')
             if model_path and os.path.exists(model_path):
-                self.model = torch.hub.load(model_path, model_version, pretrained=pretrained, source='local', device_map = self.device)
+                self.model = torch.hub.load(model_path, model_version, pretrained=pretrained, source='local')
             else:
-                self.model = torch.hub.load('ultralytics/yolov5', model_version, pretrained=pretrained, device_map = self.device)
+                self.model = torch.hub.load('ultralytics/yolov5', model_version, pretrained=pretrained)
         except Exception as e:
             print(f"Error loading YOLOv5 model: {e}")
             raise