diff --git a/utils/torch_utils.py b/utils/torch_utils.py index cac1abb4..6d10004c 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -3,21 +3,22 @@ import torch def init_seeds(seed=0): torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) def select_device(force_cpu=False): if force_cpu: + cuda = False device = torch.device('cpu') else: - device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') + 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. Multi-GPU issue: https://github.com/ultralytics/yolov3/issues/21') + 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 ' + str(device) + '\n') + print('Using %s\n%s' % (device.type, torch.cuda.get_device_properties(0) if cuda else '')) return device