To make W0 trainable, you would have to define it as an nn.Parameter.
Also, it should be used in the computation of the output of your model.
Currently it seems you are only using W0 in your __init__ method to calculate self.h, which is then overwritten in the forward method, so I’m not sure how the computation graph should look like.