[런타임에러] RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

모델(weight tensor)과 입력 데이터(input tensor) 중 하나라도 GPU에 올라가있지 않을 때 생기는 오류라고 한다. 즉, CPU 사용할거면 모두 CPU를, GPU 사용할거면 모두 GPU를 사용해야 한다.

http://education.molssi.org/gpu_programming_beginner/aio/index.html

 

수정 전

#수정 전: input data만 GPU에 올라가있음
network = ResNet50(input_channel=3, outputs=num_classes)

images = Variable(images).cuda()
labels = Variable(labels).cuda()

 

수정 후

network = ResNet50(input_channel=3, outputs=num_classes).cuda()

images = Variable(images).cuda()
labels = Variable(labels).cuda()

model.cuda()로 모델의 weight도 gpu 연산 가능하게 만드니까 잘 동작했다.

 

하지만 위 방법은 hardcoding 방식이고, GPU 사용 가능한지 체크해서 연산 device를 자동으로 설정해줄 수도 있다. 아래 general 코드를 사용하면 좀 더 편하게 여러 환경에서 사용 가능하다.

 

General Code

dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(dev)
data = data.to(dev)

 

 

참고

https://stackoverflow.com/questions/59013109/runtimeerror-input-type-torch-floattensor-and-weight-type-torch-cuda-floatte

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

 

http://education.molssi.org/gpu_programming_beginner/aio/index.html

Fundamentals of Heterogeneous Parallel Programming with CUDA C/C++