How to add dynamically defined module and its optimizer to the lightning module?
❓ Questions and Help
What is your question?
I have a module whose configuration is defined by one forward pass, e.g. a fc layer whose input size depends on some random features inside a network. I also have to update the parameters of this module. What's the recommend way to do this?
Code
PatchSampleF is a kind of this module,
What have you tried?
I try to achieve this by defining the module and its optimizer inside the first validation step ( pl.Trainer(num_sanity_val_steps=1) ) before training. However, in the following training step, all parameters of the newly-defined module do not requires grad anymore. A workaround is to manually require_grad for these parameters. I don't know what triggered the change but I think if there is a way to register the newly-defined module and corresponding optimizer to the training process, this problem can be solved gracefully.
What's your environment?

OS: Arch linux
Packaging pip
Version 1.0.5