Device_map="auto" with error: Expected all tensors to be on the same device

Beginners
3

Any update on this? Running into the same error. I’m using gpt-j with auto mapping. The inference works with model.generate but with model(input_ids) it runs into error.

from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-j-6B",
                device_map="auto",
                load_in_8bit=False)

text = "this is a test.."
input_ids = tokenizer(text, return_tensors="pt").input_ids.to("cuda")

with torch.no_grad():
    target_ids = input_ids.clone()
    full_outputs = model(input_ids, labels=target_ids)