Solved by mv15885 in post #5 
              
              
                It is indeed not right as self.U should be parameterized by self.weights, however I have managed to get it working, thanks for the help! 
I moved self.U = self.Rx().cdouble() from init to forward in the Rx_layer class. This makes sense as it now generates U with updated parameters when required.