[FIXED] "can't optimize a non-leaf Tensor" on torch Parameter

Issue

I'm getting "can't optimize a non-leaf Tensor" on this bit of code

self.W_ch1 = nn.Parameter(
            torch.rand(encoder_feature_dim, encoder_feature_dim), requires_grad=True
        ).to(self.device)
self.W_ch1_optimizer = torch.optim.Adam([self.W_ch1], lr=encoder_lr)

Don't know why it's happening that should be the leaf tensor, because it has no children connected to it. It's just a torch.rand inside a nn.Parameter variable. It throws the error at the initialization of self.w_ch1_optmizer

Solution

The reason why it throws an error is that torch.Tensor.cuda has the effect of creating a reference for transferring the data doing so by registering a new node in the graph. In other words your parameter module W_ch1 is no longer a leaf node since you already have this "computation" tree:

nn.Parameter -> cuda:parameter = W_ch1

You can compare the following two results:

>>> p = nn.Parameter(torch.rand(1)).cuda()
>>> p.is_leaf
False

What you need to be doing is first instantiate your modules, and define your optimizer(s). Only then can you transfer them to the desired device. Not before:

>>> p = nn.Parameter(torch.rand(1))
>>> optimizer = optim.Adam([p], lr=lr)

Then you can transfer everything:

>>> p.cuda()
>>> optimizer.cuda()

Answered By - Ivan