Error caused by device_map="auto"

Dear community,

toxicity_model = AutoModelForSequenceClassification.from_pretrained(toxicity_model_name, device_map=“auto”)

When I include device_map=“auto” shown above, I got the following error. Codes ran fine without device_map=“auto”. Any idea why?

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)

====================
Original code snippets:-

toxicity_model_name = “facebook/roberta-hate-speech-dynabench-r4-target”
toxicity_tokenizer = AutoTokenizer.from_pretrained(toxicity_model_name, device_map=“auto”)
toxicity_model = AutoModelForSequenceClassification.from_pretrained(toxicity_model_name, device_map=“auto”)

non_toxic_text = “#Person 1# tells Tommy that he didn’t like the movie.”

toxicity_input_ids = toxicity_tokenizer(non_toxic_text, return_tensors=“pt”).input_ids

logits = toxicity_model(input_ids=toxicity_input_ids).logits
print(f’logits [not hate, hate]: {logits.tolist()[0]}')

Print the probabilities for [not hate, hate]

probabilities = logits.softmax(dim=-1).tolist()[0]
print(f’probabilities [not hate, hate]: {probabilities}')

get the logits for “not hate” - this is the reward!

not_hate_index = 0
nothate_reward = (logits[:, not_hate_index]).tolist()
print(f’reward (high): {nothate_reward}')

I met the same problem. It seems that ‘‘toxicity_input_ids’’ is on cpu. I add “to.(cuda)” after “toxicity_input_ids = toxicity_tokenizer(non_toxic_text, return_tensors=“pt”).input_ids” and it works for me.