为什么在前向函数中定义的张量不能自动转换为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

设计使然。

只有属于模型一部分的张量会随着store = [] 或store.append(i)移动。

这些张量已向model.cuda()或model.to("cuda")注册。这还包括向上述功能注册的子模块，参数和缓冲区。

尽管register_parameter实际上是类本身的一部分，但根据设计，它不会被移至CUDA，相反，如果您没有使用{{ 1}}方法。

Answer 2

一种解决方案：

device=whatever_parameter_in_your_forward_function_args.device
self.a=torch.zeros(1).to(device)