File size: 603 Bytes
3b2b066 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
import torch
def _prepare_device(n_gpu_use):
n_gpu = torch.cuda.device_count()
if n_gpu_use > 0 and n_gpu == 0:
print("Warning: There\'s no GPU available on this machine," "training will be performed on CPU.")
n_gpu_use = 0
if n_gpu_use > n_gpu:
print("Warning: The number of GPU\'s configured to use is {}, but only {} are available " "on this machine.".
format(n_gpu_use, n_gpu))
n_gpu_use = n_gpu
device = torch.device('cuda:0' if n_gpu_use > 0 else 'cpu')
list_ids = list(range(n_gpu_use))
return device, list_ids |