updates
This commit is contained in:
parent
e28437720d
commit
e1425b7288
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue