Skip to content

Commit ac89e0a

Browse files
changed code for Single GPU only
1 parent 067d25d commit ac89e0a

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

detection_helpers.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ def __init__(self, conf_thres:float = 0.25, iou_thresh:float = 0.45, agnostic_nm
2020
save_conf: whether to save confidences in 'save_txt' labels afters inference
2121
classes: Filter by class from COCO. can be in the format [0] or [0,1,2] etc
2222
'''
23-
self.device = select_device("cuda" if torch.cuda.is_available() else 'cpu')
23+
self.device = select_device("0" if torch.cuda.is_available() else 'cpu')
2424
self.conf_thres = conf_thres
2525
self.iou_thres = iou_thresh
2626
self.classes = classes

0 commit comments

Comments
 (0)