Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-07-16 19:09:40 +02:00
parent ab51380448
commit a773350224
1 changed files with 7 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
utils/torch_utils.py
View File

@ -18,7 +18,7 @@ def select_device(force_cpu=False):
if cuda: if cuda:
try: # Mixed precision training https://github.com/NVIDIA/apex try: # Mixed precision training https://github.com/NVIDIA/apex
from apex import amp from apex import amp
apex_str = ' with Apex' apex_str = 'with Apex '
except: except:
apex_str = '' apex_str = ''
@ -27,13 +27,12 @@ def select_device(force_cpu=False):
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_str cuda_str = 'Using CUDA ' + apex_str
print("%sdevice0 _CudaDeviceProperties(name='%s', total_memory=%dMB)" % for i in range(0, ng):
(cuda_str, x[0].name, x[0].total_memory / c)) if ng == 1:
if ng > 0: # torch.cuda.set_device(0) # OPTIONAL: Set GPU ID
# torch.cuda.set_device(0) # OPTIONAL: Set GPU ID cuda_str = ' ' * len(cuda_str)
for i in range(1, ng): print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" %
print("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" % (cuda_str, i, x[i].name, x[i].total_memory / c))
(' ' * len(cuda_str), i, x[i].name, x[i].total_memory / c))
print('') # skip a line print('') # skip a line
return device return device