已修复Tensor for argument #2 ‘weight‘ is on CPU, but expected it to be on GPU

Posted 2022-02-25 MarToony|名角

Bug提示信息

RuntimeError: Tensor for argument #2 'weight' is on CPU, but expected it to be on GPU (while checking arguments for cudnn_batch_norm).

深层原因

参考文献 RuntimeError：参数#2“权重”的张量在 CPU 上，但预计它在 GPU 上（同时检查 cudnn_batch_norm 的参数）
是因为，代码中，正在动态地在forward方法内创建一些层。
这些层将在每次前向传递中重新初始化，因此它们不会被训练，也不会被推送到 GPU。
标准方法是在方法中创建这些层__init__（与您的其他层一样），然后将它们应用于forward.

直接原因：

深层原因过于抽象，放到我的代码环境中来看，是因为列表表达式生成的N个网络层，没有被nn.ModuleList方法包裹。这是pytorch的设计要求。

self.initial_stage = nn.ModuleList([Block(
    init_channel=stem_channel, input_channel=num_channel, **kwargs) for _ in range(num_input)])

具体来说，上面的Block层中含有Batchnorm网络层，self.initial_stage在定义的时候没有被nn.ModuleList所包裹，使得在forward中对self.initial_stage进行调用的时候，触发了上面的Bug。