python 自動獲取可用GPU的ID

from gpuinfo import GPUInfo

def get_gpu_id(num_gpu=1):
    """get ID of GPUS
    :param num_gpu:  num of GPUs to use
    :return: gpu_id: ID of allocated GPUs
    """
    available_device=GPUInfo.check_empty()
    if len(available_device)>=num_gpu:
        gpu_id = available_device[:num_gpu]
    else:
        raise Exception('Only {} GPUs to use!'.format(len(available_device)))
    if num_gpu == 1:
        gpu_id = gpu_id[0]


    return gpu_id