diff --git a/utils/torch_utils.py b/utils/torch_utils.py index a45be2c1..c1629f3a 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -15,10 +15,10 @@ def select_device(force_cpu=False): cuda = torch.cuda.is_available() device = torch.device('cuda:0' if cuda else 'cpu') - if torch.cuda.device_count() > 1: - print('WARNING Using GPU0 Only: https://github.com/ultralytics/yolov3/issues/21') - torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available - # print('Using ', torch.cuda.device_count(), ' GPUs') + # if torch.cuda.device_count() > 1: + # print('WARNING Using GPU0 Only: https://github.com/ultralytics/yolov3/issues/21') + # torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available + # # print('Using ', torch.cuda.device_count(), ' GPUs') print('Using %s %s\n' % (device.type, torch.cuda.get_device_properties(0) if cuda else '')) return device