diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 8219ccc9..2fa2b275 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -21,7 +21,7 @@ def select_device(force_cpu=False, apex=False): c = 1024 ** 2 # bytes to MB ng = torch.cuda.device_count() x = [torch.cuda.get_device_properties(i) for i in range(ng)] - cuda_str = 'Using CUDA ' + 'Apex ' if apex else '' + cuda_str = 'Using CUDA ' + ('Apex ' if apex else '') for i in range(0, ng): if i == 1: # torch.cuda.set_device(0) # OPTIONAL: Set GPU ID