PyTorch DataParallel에서 기본 GPU와 지정 GPU 불일치 문제 해결 방법

gpu가 8개인 서버에서 5,6,7,8번 째 gpu에만 데이터 병렬처리를 하려고 하는데 아래 오류가 있었다.
RuntimeError: module must have its parameters and buffers on device cuda:4 (device_ids[0]) but found one of them on device: cuda:0  

 

model = nn.DataParallel(model, device_ids=gpus).cuda()

gpus는 (4,5,6,7) 이었다. 

 

문제는 모델을 gpus[0]으로 옮긴 후 장치를 지정하지 않고 .cuda()를 호출하면 모델이 기본 GPU(일반적으로 cuda:0)로 전송되어 지정한 기본 장치와 불일치가 발생한다는 것입니다. 이를 수정하려면 DataParallel로 래핑한 후 추가 .cuda() 호출을 제거하거나 gpus[0]으로 다시 명시적으로 이동하기만 하면 됩니다. 예를 들어 다음과 같이 코드를 업데이트합니다.

 

model = nn.DataParallel(model, device_ids=gpus).to(gpus[0])