모델 학습 GPU 병행처리(nn.DataParallel)

두 개이상의 GPU가 있다면 여러개의 GPU 사용이 가능하다.

import torch
import torch.nn as nn

if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"
    
 "사용할 모델" = 미리 구성
 
if torch.cuda.device_count() > 1:
    model = nn.DataParallel("사용할 모델").to(device)