[Pytorch] Multi GPU

pytorch에서 여러개의 GPU 사용하기.

 

import torch

model = MyModel()          # CNN이든 뭐든 사용할 모델
device = 'cuda' if torch.cuda.is_available() else 'cpu'

if (device == 'cuda') and (torch.cuda.device_count() > 1):
    model = nn.DataParallel(model)

model.to(device)

 

간단하다.