RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0!

#3
by tiansiyuan - opened

Running the example code:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper__index_select)

Adding "device_map='auto'" solves it.

model = AutoModelForCausalLM.from_pretrained("codefuse-ai/CodeFuse-CodeLlama-34B", trust_remote_code=True, device_map='auto')

tiansiyuan changed discussion status to closed

Sign up or log in to comment