[Pytorch] RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

pytorch tensorboard tutorial를 참조하며 코드를 작성한 후 모델을 돌려보던 중 다음과 같은 에러를 마주하였다.

 

 

pytorch tensorboard tutorial

https://tutorials.pytorch.kr/intermediate/tensorboard_tutorial.html

TensorBoard로 모델, 데이터, 학습 시각화하기 — PyTorch Tutorials 1.10.0+cu102 documentation

 

대략적인 코드

# 자세한 코드 내용은 생략

# net은 torch로 쌓은 neural network임
model = net().cuda()

# trainloader를 거쳐 데이터 하나를 뽑은 것이 images임
dataiter = iter(trainloader)
images, labels = dataiter.next()

from torch.utils.tensorboard import SummaryWriter

# tensorboard
writer = SummaryWriter('logs')

# 에러가 난 부분
writer.add_graph(model, images)
writer.close()

 

마주친 에러

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

 

원인

이는 model은 cuda 처리를 해주었으나 image에는 cuda처리를 해주지 않아 발생한 에러이다.

 

 

해결방법

다음과 같이 바꿔주면

 

# 수정한 부분
writer.add_graph(model, images.cuda())
writer.close()

 

잘 돌아간다. 

 

 

 

-히비스서커스-