Search code examples
pytorchcudagpuautograd

Assigning a parameter to the GPU sets is_leaf as false

If I create a Parameter in PyTorch, then it is automatically assigned as a leaf variable:

x = torch.nn.Parameter(torch.Tensor([0.1]))
print(x.is_leaf)

This prints out True. From what I understand, if x is a leaf variable, then it will be updated by the optimiser.

But if I then assign x to the GPU:

x = torch.nn.Parameter(torch.Tensor([0.1]))
x = x.cuda()
print(x.is_leaf)

This prints out False. So now I cannot assign x to the GPU and keep it as a leaf node.

Why does this happen?

Solution

  • Answer is in is_leaf documentation and here is your exact case:

    >>> b = torch.rand(10, requires_grad=True).cuda()
>>> b.is_leaf
False
# b was created by the operation that cast a cpu Tensor into a cuda Tensor

    Citing documentation further:

    For Tensors that have requires_grad which is True, they will be leaf Tensors if they were created by the user. This means that they are not the result of an operation and so grad_fn is None.

    In your case, Tensor was not created by you, but was created by PyTorch's cuda() operation (leaf is the pre-cuda b).