使用nn.LayerNorm报错found at least two devices

原本的代码是这个：

import torch
import torch.nn as nn

N, C, H, W = 20, 5, 10, 10
input = torch.randn(N, C, H, W).to('cuda')
layer_norm = nn.LayerNorm([C, H, W])
output = layer_norm(input)
'''
RuntimeError: Expected all tensors to be on the same device,
 but found at least two devices, cuda:0 and cpu! 
(when checking argument for argument weight in method wrapper__native_layer_norm)
'''

这里的话，应该主要就是网络没有送到GPU里面，所以会报错，所以修改，将归一化部分网络送入GPU即可：

import torch.nn as nn

N, C, H, W = 20, 5, 10, 10
input = torch.randn(N, C, H, W).to('cuda')
layer_norm = nn.LayerNorm([C, H, W]).to('cuda')
output = layer_norm(input)