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

여태까지 PyTorch 사용한 방식과 차이가 없었던 거 같은데 저런 문구가 떠서 조금 당황스러웠다. BERT 실습 중에 만난 거였고 일단 해결하긴 했는데 깔끔한 방식은 아닌 것 같다.

 

 

빨간색 체크 표시를 해둔 곳에서 RuntimeError: ~ 가 떴고 모든 텐서가 cuda:0에 올라가야 하는데 cpu에 올라간 게 있는 것 같아서 빨간색 체크 전에 print()로 각 텐서가 어디 올라가있는지 device를 찍어보았다.

 

model만 GPU에 올라가있는 것 확인

 

input_ids, attention_mask, token_type_ids는 custom torch dataset, collate_fn, dataloader를 거쳐 최종적으로 모델 학습 데이터가 되기 때문에 collate_fn의 리턴 결과에 모두 .to(device)를 추가해주었다.

 

모두 GPU에 올라간 것을 확인했다.