def select_device(device='', batch_size=None):
# device = 'cpu' or '0' or '0,1,2,3'
cpu = device.lower() == 'cpu'
if cpu:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
elif device: # non-cpu device requested
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
cuda = not cpu and torch.cuda.is_available()
return torch.device('cuda:0' if cuda else 'cpu')
显卡编码问题可以参考这篇文章pytorch之多GPU使用