@smth,  your solution assumes we are defining a custom neural network module where W is a parameter of the network, right?
Is there any way to tie weights by only instantiating off-the-shelf nn modules without fiddling around with new parameters?