added initializers and self.layers for resetting weights

osl_dynamics/inference/layers.py

@@ -913,17 +913,14 @@ def __init__(
         self.damping = LearnableTensorLayer(
             shape=(n, 1),
             learn=learn,
-            initializer=None,
-            initial_value=0.5 * tf.ones((n, 1)),
+            initializer=initializers.Constant(0.5),
             name=self.name + "_damping",
         )
 
         self.frequency = LearnableTensorLayer(
             shape=(n, 1),
             learn=learn,
-            initializer=None,
-            initial_value=tf.random.uniform(
-                (n, 1),
+            initializer=initializers.RandomUniform(
                 minval=frequency_limit[0],
                 maxval=frequency_limit[1],
             ),
@@ -933,11 +930,12 @@ def __init__(
         self.amplitude = LearnableTensorLayer(
             shape=(n, 1),
             learn=learn and learn_amplitude,
-            initializer=None,
-            initial_value=tf.ones((n, 1)),
+            initializer=initializers.Constant(1.0),
             name=self.name + "_amplitude",
         )
 
+        self.layers = [self.damping, self.frequency, self.amplitude]
+
     def call(self, inputs, **kwargs):
         """Calculate damped oscillator.
 
@@ -1003,6 +1001,8 @@ def __init__(
             learn=learn,
         )
 
+        self.layers = [self.oscillator_layer]
+
     def call(self, inputs, **kwargs):
         """Retrieve the covariance matrices.