HemaAM commited on
Commit
0ce0c39
·
1 Parent(s): d4fb5fa

Bug fix to not use cuda

Browse files
Files changed (1) hide show
  1. 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)]