Spaces:
Sleeping
Sleeping
Bug Fix for use_cuda
Browse files- 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=
|
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()
|