This commit is contained in:
Glenn Jocher 2019-07-24 18:30:35 +02:00
parent e28437720d
commit e1425b7288
1 changed files with 1 additions and 1 deletions

View File

@ -21,7 +21,7 @@ def select_device(force_cpu=False, apex=False):
c = 1024 ** 2 # bytes to MB c = 1024 ** 2 # bytes to MB
ng = torch.cuda.device_count() ng = torch.cuda.device_count()
x = [torch.cuda.get_device_properties(i) for i in range(ng)] 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): for i in range(0, ng):
if i == 1: if i == 1:
# torch.cuda.set_device(0) # OPTIONAL: Set GPU ID # torch.cuda.set_device(0) # OPTIONAL: Set GPU ID