Spaces:
Sleeping
Sleeping
Bug fix to not use cuda
Browse files- detection.py +1 -1
detection.py
CHANGED
@@ -26,7 +26,7 @@ yolo_grad_cam = YoloGradCAM(model=model, target_layers=[model.layers[-2]], use_c
|
|
26 |
@torch.inference_mode()
|
27 |
def detect_objects(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.4, enable_grad_cam: bool = False, transparency: float = 0.5) -> List[np.ndarray]:
|
28 |
transformed_image = modelConfig.transforms(image=image)["image"].unsqueeze(0)
|
29 |
-
transformed_image = transformed_image.cuda()
|
30 |
output = model(transformed_image)
|
31 |
|
32 |
bboxes = [[] for _ in range(1)]
|
|
|
26 |
@torch.inference_mode()
|
27 |
def detect_objects(image: np.ndarray, iou_thresh: float = 0.5, thresh: float = 0.4, enable_grad_cam: bool = False, transparency: float = 0.5) -> List[np.ndarray]:
|
28 |
transformed_image = modelConfig.transforms(image=image)["image"].unsqueeze(0)
|
29 |
+
#transformed_image = transformed_image.cuda()
|
30 |
output = model(transformed_image)
|
31 |
|
32 |
bboxes = [[] for _ in range(1)]
|