HemaAM commited on
Commit
d4fb5fa
·
1 Parent(s): 5c5fbf4

Bug Fix for use_cuda

Browse files
Files changed (1) hide show
  1. detection.py +1 -1
detection.py CHANGED
@@ -20,7 +20,7 @@ scaled_anchors = (
20
  * torch.tensor(modelConfig.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
21
  ).to(modelConfig.DEVICE)
22
 
23
- yolo_grad_cam = YoloGradCAM(model=model, target_layers=[model.layers[-2]], use_cuda=True)
24
 
25
 
26
  @torch.inference_mode()
 
20
  * torch.tensor(modelConfig.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
21
  ).to(modelConfig.DEVICE)
22
 
23
+ yolo_grad_cam = YoloGradCAM(model=model, target_layers=[model.layers[-2]], use_cuda=False)
24
 
25
 
26
  @torch.inference_mode()