python - 为什么我在 forward 函数中定义的张量不能自动转换为 cuda 变量？

Question

在 PyTorch 中，在模型的 forward 函数中

class my_model(torch.nn.module):

    ......

    def forward():
        self.a=torch.zeros(1)
        #blabla

在model.cuda()之后，为什么self.a还是一个cpu变量？

Answer 1

这是设计使然。

只有作为模型一部分的张量才会随 model.cuda() 或 model.to("cuda") 移动。

这些张量使用 register_parameter 或 register_buffer 注册。这还包括使用上述函数注册的子模块、参数和缓冲区。

尽管 self.a=torch.zeros(1) 实际上是类本身的一部分，但根据设计它不会被移动到 CUDA，相反你需要做 a.to("cuda")，如果您还没有使用 register* 方法。