[PyTorch]RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! 해결법

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

다음과 같은 에러가 발생했다.

model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2",device_map="auto", torch_dtype=torch.float16)

# 코드 생략

outputs = model.generate(**inputs,max_length=2048)

다음과 같은 오류가 발생했을 때, use_cache=True를 사용하면 해결할 수 있다.

model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2",device_map="auto", torch_dtype=torch.float16)

# 코드 생략

outputs = model.generate(**inputs,use_cache=True, max_length=2048)

바뀐 소스 코드

끝!