Depending on your setup, the cleanest way might be to code the module in a way that you can use register_buffer similar to BatchNorm and friends do for running statistics.
If that is not an option, I think the cleanest interface would be to pass in the Module’s __call__ method and store and call that.
Best regards
Thomas