Skip to content

Commit e82d256

Browse files
Merge pull request #19 from deshwalmahesh/cuda_issue
changed code for Single GPU only
2 parents 067d25d + ac89e0a commit e82d256

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)