TorchScript로 저장한 모델을 CPU로 불러와야 할 때.

서기현·2024년 4월 6일

DeepLearning

목록 보기
2/6
model = torch.jit.load(model_file, map_location=torch.device('cpu'))

종종 gpu로 학습한 모델을 torchscript로 저장 후, cpu 로컬 환경에서 테스트할 일이 있는데, 이때는 모델을 로드할 때 위 같이 map_location을 지정해주면 된다.

0개의 댓글