Device error for JSON generation with outlines 0.0.29

[streitl]

Describe the issue as clearly as possible:

I'm trying to use outlines to generate json objects based on a Pydantic model.
I am loading my model on a GPU, but when I try to generate an object, I get an error saying that pytorch also found tensors on the CPU.

Steps/code to reproduce the bug:

import outlines
import torch

device: str = "cuda"
rng = torch.Generator(device=device)
rng.manual_seed(90302)

outlines_model = outlines.models.transformers(model_id, device=device)
generator = outlines.generate.json(outlines_model, MyPydanticModel, whitespace_pattern="")

outlines_output = generator("test prompt", rng=rng)

Expected result:

A MyPydanticModel instance.

Error message:

Traceback (most recent call last):
  File "/my_example.py", line ..., in <module>
    outlines_output = generator("test prompt", rng=rng)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.11/site-packages/outlines/generate/api.py", line 205, in __call__
    last_state = next(states)
                 ^^^^^^^^^^^^
  File "lib/python3.11/site-packages/outlines/generate/generator.py", line 83, in sequence_generator
    token_ids = update_token_ids(token_ids, next_token_ids, ancestors)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.11/site-packages/outlines/generate/generator.py", line 198, in update_token_ids
    token_ids = torch.index_select(token_ids, 0, ancestors)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)

Outlines/Python version information:

outlines 0.0.29
Python 3.11.7 (main, Dec  8 2023, 18:56:58) [GCC 11.4.0]
accelerate==0.26.1
aiofiles==23.2.1
aiohttp==3.9.3
aiosignal==1.3.1
altair==5.2.0
annotated-types==0.6.0
anyio==4.2.0
asttokens==2.4.1
attrs==23.2.0
auto-gptq==0.6.0
bitsandbytes==0.42.0
black==23.12.1
certifi==2024.2.2
charset-normalizer==3.3.2
click==8.1.7
cloudpickle==3.0.0
coloredlogs==15.0.1
contourpy==1.2.0
cycler==0.12.1
datasets==2.16.1
decorator==5.1.1
dill==0.3.7
diskcache==5.6.3
docstring-parser==0.15
einops==0.7.0
executing==2.0.1
fastapi==0.109.2
ffmpy==0.3.1
filelock==3.13.1
flake8==6.1.0
fonttools==4.48.1
frozenlist==1.4.1
fsspec==2023.10.0
gekko==1.0.6
gradio==3.50.2
gradio_client==0.6.1
h11==0.14.0
httpcore==1.0.2
httpx==0.26.0
huggingface-hub==0.20.3
humanfriendly==10.0
idna==3.6
importlib-resources==6.1.1
interegular==0.3.3
ipdb==0.13.13
ipython==8.21.0
isort==5.13.2
jedi==0.19.1
jieba==0.42.1
Jinja2==3.1.3
joblib==1.3.2
jsonschema==4.21.1
jsonschema-specifications==2023.12.1
kiwisolver==1.4.5
lark==1.1.9
llama_cpp_python==0.2.39
llmtuner==0.5.1
llvmlite==0.42.0
markdown-it-py==3.0.0
MarkupSafe==2.1.5
matplotlib==3.8.2
matplotlib-inline==0.1.6
mccabe==0.7.0
mdurl==0.1.2
mpmath==1.3.0
multidict==6.0.5
multiprocess==0.70.15
mypy==1.8.0
mypy-extensions==1.0.0
nest-asyncio==1.6.0
networkx==3.2.1
nltk==3.8.1
numba==0.59.0
numpy==1.26.4
nvidia-cublas-cu12==12.1.3.1
nvidia-cuda-cupti-cu12==12.1.105
nvidia-cuda-nvrtc-cu12==12.1.105
nvidia-cuda-runtime-cu12==12.1.105
nvidia-cudnn-cu12==8.9.2.26
nvidia-cufft-cu12==11.0.2.54
nvidia-curand-cu12==10.3.2.106
nvidia-cusolver-cu12==11.4.5.107
nvidia-cusparse-cu12==12.1.0.106
nvidia-nccl-cu12==2.19.3
nvidia-nvjitlink-cu12==12.3.101
nvidia-nvtx-cu12==12.1.105
optimum==1.16.2
orjson==3.9.13
outlines==0.0.29
packaging==23.2
pandas==2.2.0
pandas-stubs==2.1.4.231227
parso==0.8.3
pathspec==0.12.1
peft==0.8.2
pexpect==4.9.0
pillow==10.2.0
platformdirs==4.2.0
prompt-toolkit==3.0.43
protobuf==4.25.2
psutil==5.9.8
psycopg==3.1.18
psycopg-binary==3.1.18
ptyprocess==0.7.0
pure-eval==0.2.2
pyarrow==15.0.0
pyarrow-hotfix==0.6
pycodestyle==2.11.1
pydantic==2.6.1
pydantic-2-mermaid==0.6.0
pydantic_core==2.16.2
pydub==0.25.1
pyflakes==3.1.0
Pygments==2.17.2
pyparsing==3.1.1
python-dateutil==2.8.2
python-multipart==0.0.7
pytz==2024.1
PyYAML==6.0.1
referencing==0.33.0
regex==2023.12.25
requests==2.31.0
rich==13.7.0
rouge==1.0.1
rouge-chinese==1.0.3
rpds-py==0.17.1
safetensors==0.4.2
scikit-learn==1.4.0
scipy==1.12.0
semantic-version==2.10.0
sentence-transformers==2.3.1
sentencepiece==0.1.99
shtab==1.6.5
six==1.16.0
sniffio==1.3.0
sse-starlette==2.0.0
stack-data==0.6.3
starlette==0.36.3
sympy==1.12
threadpoolctl==3.2.0
tokenizers==0.15.1
toolz==0.12.1
torch==2.2.0
tqdm==4.66.1
traitlets==5.14.1
transformers==4.37.2
triton==2.2.0
trl==0.7.10
types-pytz==2024.1.0.20240203
types-tqdm==4.66.0.20240106
typing_extensions==4.9.0
tyro==0.7.2
tzdata==2023.4
urllib3==2.2.0
uvicorn==0.27.0.post1
wcwidth==0.2.13
websockets==11.0.3
xxhash==3.4.1
yarl==1.9.4

Context for the issue:

This used to work in version 0.0.27, I think this might be related to 3391fa4

[lapp0]

Could you please try #651

pip install git+https://github.com/lapp0/outlines@ensure-ancestors-on-same-device

[rlouf]

The fix was just merged on main, I’m cutting a new release

[streitl]

Thank you for the quick fix, I tested lapp0's branch and it works well! 🎉