If I make a new module to use in a NN, should I always be adding the @weak_module annotation, or in general, when would I want to use this for my own modules?
Specifically in my case, I was considering subclassing torch.nn.conv2d so that I could manipulate the weights before sending it through the normal conv2d module, so should I add @weak_module to my new subclass?
The original conv2d source has it, so wasn’t sure if my modules also need it.
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)
But would using @weak_module and @weak_script make my custom nn changes faster, or no?
@weak_module and @weak_script only delay compilation until the module they’re attached to is used, so there would be no performance change with it.