import torch


def select_device(gpu):
    """
    选择CPU或GPU
    Args:
        gpu (bool): 选择GPU如果True，否则CPU

    Returns:
        device (torch.device): 'cpu' or 'cuda:0'
    """
    if gpu:
        assert torch.cuda.is_available(), ('Error: Requested GPU, '
                                           'but GPU is not available.')

    #选择device
    device = torch.device('cuda:0') if gpu else torch.device('cpu')
    return device