I don’t get it at all but since you are adding new layers I imagine you have to pass them to the optimizer at some point right?
Soo the leaf tensor is the original nn.Module. Once you allocate it in cuda, the output of that allocation is a non-leaf tensor whose backward is something like “copy to the cpu leaf node”
Soooo I would tell you to pass the cpu module to the optimizer or to convert the one in cuda into a leaf tensor.
Anyway it would be nice if you paste a code to reproduce it.