import torch


def gpu_device(gpu):
    """
    Returns a device based on the passed parameters.

    Parameters
    ----------
    gpu: bool or int
        If int, the returned device is the GPU with the specified ID. If False, the returned device
        is the CPU, if True, the returned device is given as the GPU with the highest amount of
        free memory.

    Returns
    -------
    torch.device
        A PyTorch device.
    """
    if isinstance(gpu, bool) and gpu:
        assert torch.cuda.is_available()
        return torch.device('cuda', 0)

    if isinstance(gpu, bool):
        return torch.device('cpu')

    assert gpu < torch.cuda.device_count()
    return torch.device('cuda', gpu)