The most elegant way would probably be the functional API, which would only create a single weight parameter and just use it when it’s needed.
Alternatively, you could assign the weight parameter to your modules as described here.
1 Like