Hi,
First you should never use .data.
When you do model.mu.data it returns a new Tensor and so its .grad field will always be None.
Also, you want to move     self.center_freq = (self.freq_scale*self.mu.repeat(1,self.K[0])) to the forward funciton:  every computation that is part of the forward should happen in the forward.
Otherwise, you will have part of these computations that are shared across iterations and it will lead to errors.