모델 저장
import torch
torch.save(model.state_dict(), PATH)
모델 불러오기
GPU에서 save, CPU에서 load
import torch
device = torch.device('cpu')
model = CNN(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
GPU에서 save, GPU에서 load
import torch
device = torch.device('cuda')
model = CNN(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
CPU에서 save, GPU에서 load
import torch
device = torch.device('cuda')
model = CNN(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))
model.to(device)