Hi,
The only gradients you actually care about are for weight_real no?
In that case you can just make weight_real a Parameter of your module and create weight in a differentiable manner during the forward.