You shouldn’t need to worry about the @weak_module annotation, that is internal to JIT to make the nn library compatible and low overhead with JIT, it is not a public API so it means use it at your own risk.
In practice, you don’t need to care about that annotation in most cases. If you want to extend a nn class and use it in TorchScript, I encourage you to make a ScriptModule as indicated by the JIT documentation (https://pytorch.org/docs/stable/jit.html 12)