Skip to content

Commit

Permalink
Merge pull request #19 from deshwalmahesh/cuda_issue
Browse files Browse the repository at this point in the history
changed code for Single GPU only
  • Loading branch information
deshwalmahesh authored Dec 30, 2022
2 parents 067d25d + ac89e0a commit e82d256
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion detection_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ def __init__(self, conf_thres:float = 0.25, iou_thresh:float = 0.45, agnostic_nm
save_conf: whether to save confidences in 'save_txt' labels afters inference
classes: Filter by class from COCO. can be in the format [0] or [0,1,2] etc
'''
self.device = select_device("cuda" if torch.cuda.is_available() else 'cpu')
self.device = select_device("0" if torch.cuda.is_available() else 'cpu')
self.conf_thres = conf_thres
self.iou_thres = iou_thresh
self.classes = classes
Expand Down

0 comments on commit e82d256

Please sign in to comment.