nn.dataparallel 사용 시 유의할 점

GPU 4개가 장착되어 있는 환경에서 train #1은 GPU #0, #1에, train #2는 GPU #2, #3에 올려서 사용하고 싶었다.

 

model = nn.DataParallel(model, device_ids = [2,3], output_device=2).to(device)

이와 같이 코드 작성시 device 0에 값들이 있다면서 에러가 발생하였다.

 

device = "cuda" if torch.cuda.is_available() else "cpu"

이 경우 위와 같이 설정되있었던 device 역시 아래와 같이 특정 GPU로 지정을 해주어야 해당 GPU를 기본으로 사용하게된다. 

device = "cuda:2" if torch.cuda.is_available() else "cpu"