import torch


def device(use_cuda):
    """
    setup GPU device if available, move model into configured device
    # 如果n_gpu_use为数字，则使用range生成list
    # 如果输入的是一个list，则默认使用list[0]作为controller
    Example:
        use_gpu = '' : cpu
        use_gpu = '0': cuda:0
        use_gpu = '0,1' : cuda:0 and cuda:1
     """

    if not use_cuda:
        device_type = 'cpu'
    else:
        device_type = f"cuda"
    device = torch.device(device_type)
    return device