pytorch 出现RuntimeError: CUDA error: an illegal memory access was encountered的解决方法

最近在使用pytorch编写transformer模型时遇到一个奇怪的报错，使用CPU训练模型时没有问题，但是一切换到GPU训练时就出问题，会报
RuntimeError: CUDA error: an illegal memory access was encountered

1. 首先在文件头插入以下代码，可以让报错信息更准确

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

2. 之后再运行模型，发现报错在forward函数的最后一行上

return nn.LayerNorm(d_model)(output + residual)

这行代码上，经过查询资料得知，这是因为nn.LayerNorm这个layer还储存在cpu上，要计算cuda上保存的变量时就会报错。所以正确的解决方法就是在构造模型时就声明nn.LayerNorm，不要在forward中声明nn.layerNorm

class transformer(nn.Module):
  def __init__(self):
    ###代码###
    xxxxx
    ##########
    self.layerNorm = nn.LayerNorm(d_model)
  def forward(self):
    ###代码###
    xxxxx
    ##########
    return self.layerNorm(output + residual)
    

希望我的经验能帮助到你 =w=