import torch


def get_device(cfg):
    """ Get device based on configuration
    Args: 
        cfg (dict): a dict of configuration
    Returns:
        torch.device
    """
    device = None
    if cfg['device'] == 'cpu':
        device = torch.device("cpu")
    elif cfg['device'] == 'cuda:0':
        device = torch.device("cuda:0")
    elif cfg['device'] == 'cuda:1':
        device = torch.device("cuda:1")
    else:
        raise NotImplementedError
    return device