If self.n and self.d are initial states which you want to optimize as additional parameters, then they should be nn.Parameter instances and shouldn’t be assigned to during forward. If they’re just temporary variables to hold state during the forward pass, they should not be attributes of self and should be created and assigned to as ordinary local variables during forward.