Hi Everyone,
I have a question about build net in nn.Module.
I use self.heads to store some nets as heads of my net, and use for loop in forward, cause I don’t want to write another module when I need change the number of heads.
The problem is model.parameters() do not contain heads’ parameters, just backbone network. How can I do to make model.parameters() recognized the heads as submodules?
self.heads is a list whose elements is nn.Module.
I guess if I manually extracts parameters as param_group, I also can use optimizer correctly and the dynamic graph design should back prop as I wish.
But it is urgely, do you guys have some suggest? Or I am wrong, the technic doesn’t work.
Please let me know.
You can use nn.ModuleList() or nn.ModuleDict() to store the nn.Modules.
oh, thank you! My problem solved.