import torch


def get_device(logger):
    """
    Get device model will be run on (GPU or CPU)
    :param logger: Logger object to note the device
    :return: device type, num_of_gpus
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    n_gpu = torch.cuda.device_count()
    logger.info("device: {}, n_gpu {}".format(device, n_gpu))
    return device, n_gpu