[pytorch] Expected object of device type cuda but got device type cpu for argument #2 'weight' 오류를 해결하는 방법 정리

pytorch 코드를 짜다 보면 다양한 오류를 만나게 된다. shape 관련한 오류들이 제일 먼저 귀찮게 하며, 해당 오류를 다 정리하고 나면 오류의 수는 줄어들지만 더 까다로운 오류들이 등장한다. 오늘은 그 중에서도 특히 해결하기 힘들었지만 무척이나 간단한 실수였던 Expected object of device type cuda but got device type cpu for argument #2 'weight' 오류를 어떻게 하면 해결하는지 정리하여 해당 오류로 오랫동안 삽질을 하는 사람들에게 도움이 되고자 한다.

 

해당 오류의 증상

해당 오류는 pytorch에서 모델에 대한 코드를 잘 작성했고, 모델의 입력 tensor가 gpu 상에 위치하고 있음에도 불구하고 발생한다. 본인의 경우 정확히는 아래와 같은 오류가 출력되었다.

  File "/home/admin/anaconda3/envs/modelss/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 202, in forward
    self.padding, self.dilation, self.groups)
RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'weight' in call to _thnn_conv2d_forward

오류의 내용이 말하기를 모델의 weight가 gpu상에 위치하고 있지 않다고 한다. 모델의 weight를 cpu에서 gpu 상으로 옮기면 되는 것인데 pytorch의 convolution layer에 대한 문서를 아무리 찾아봐도 device를 지정하는 parameter나 함수를 찾을 수 없었다.

해당 오류를 해결하는 방법

오류를 해결하는 방법은 간단했다. 모델의 인스턴스를 생성한 뒤 모델.cuda() 함수를 실행시켜주기만 하면 된다.

...
model = MyModel()
model.cuda() # 이 코드를 추가하면 문제가 해결된다.
...

이런 간단한 함수 하나로 해결되지 않았던 문제가 금새 해결되었다.