(pytorch)model.cuda()をしたはずなのに一部cudaになってないエラー

まとめ

• model.cuda()をしてTrainしようと思ったらconvのパラメータがcudaに入ってないと言われる
• いろいろ試してみた所、model.cuda()をした後にレイヤーを変えたりしたのが原因だった
• 全部終わった後にmodel.cuda()をすることで対応

コードサンプル

エラー出る例

model = torchvision.models.resnet18(pretrained=True).cuda()
model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3,bias=False)
model.fc = nn.Linear(512 * torchvision.models.resnet.BasicBlock.expansion,1)

このままトレーニングしようとすると型違うというエラーが出る

エラー出ない例

model = torchvision.models.resnet18(pretrained=True)
model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3,bias=False)
model.fc = nn.Linear(512 * torchvision.models.resnet.BasicBlock.expansion,1)
model.cuda()

最後にまるっとcudaに突っ込むことでエラー回避