Search code examples
pythonpytorchtorch

Why my tensor defined in forward function can not be transformed in to cuda variable autonomously?

In PyTorch, in the forward function of a model 

class my_model(torch.nn.module):

    ......

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

After model.cuda(), why self.a is still a cpu variable?

Solution

  • This is so by design.

    Only the tensors which are a part of the model will move with model.cuda() or model.to("cuda").

    These tensors are registered with register_parameter or register_buffer. This also includes child modules, parameters and buffers registered with the aforementioned functions.

    Even though self.a=torch.zeros(1) is actually a part of the class itself, by design it will not be moved to CUDA, instead you would need to do a.to("cuda"), if you haven't used the register* methods.