Your approach could work theoretically, but I would be careful about the overall workflow in your training script.
Usually you would setup the model and pass all parameters to an optimizer before the first forward pass.
This won’t work properly in your approach, since not all modules are initialized yet, so you would need to perform an example forward pass before creating the optimizer.
While it might work for your project, other users (using your model) might easily run into issues, where the model doesn’t train, since the parameters were lazily initialized.