import torch


def _set_device(disable_cuda=False):
    """Set device to CPU or GPU.

    Parameters
    ----------
    disable_cuda : bool (default=False)
        Whether to use CPU instead of GPU.

    Returns
    -------
    device : torch.device object
        Device to use (CPU or GPU).
    """
    # XXX we might also want to use CUDA_VISIBLE_DEVICES if it is set
    if not disable_cuda and torch.cuda.is_available():
        device = torch.device('cuda:0')
    else:
        device = torch.device('cpu')

    return device