RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! in ch05/01_main-chapter-code/gpt_generate.py

[xfpg21421]

Bug description

gpt = GPTModel(gpt_config)
load_weights_into_gpt(gpt, params)
gpt.to(device)
gpt.eval()

tokenizer = tiktoken.get_encoding("gpt2")
torch.manual_seed(123)

token_ids = generate(
    model=gpt,
    idx=text_to_token_ids(input_prompt, tokenizer),
    max_new_tokens=25,
    context_size=gpt_config["context_length"],
    top_k=50,
    temperature=1.0
)

print("Output text:\n", token_ids_to_text(token_ids, tokenizer))`

If code above runs in GPU device, error below would appear:
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)

Since the gpt model has been put to the GPU, the idx tensor from text_to_token_ids(input_prompt, tokenizer) has not been put to the GPU. So it should be modified as text_to_token_ids(input_prompt, tokenizer).to(device)

What operating system are you using?

Linux

Where do you run your code?

Other cloud environment (AWS, Azure, GCP)

Environment

[rasbt]

Good catch, thanks for reporting this. Looks like it was missing a .to(device) for the prompt.