can provide h_/c_init to initalize rnn

USGS-R · Jun 7, 2021 · 8362981 · 8362981
1 parent 1c1641a
commit 8362981
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/river_dl/rnns.py b/river_dl/rnns.py
@@ -17,6 +17,7 @@ def __init__(
         input element to be zero
         """
         super().__init__()
+        self.hidden_size = hidden_size
         self.num_tasks = num_tasks
         self.rnn_layer = layers.LSTM(
             hidden_size,
@@ -33,6 +34,10 @@ def __init__(
 
     @tf.function
     def call(self, inputs, **kwargs):
+        batch_size = inputs.shape[0]
+        h_init = kwargs.get("h_init", tf.zeros([batch_size, self.hidden_size]))
+        c_init = kwargs.get("c_init", tf.zeros([batch_size, self.hidden_size]))
+        self.rnn_layer.reset_states(states=[h_init, c_init])
         x, h, c = self.rnn_layer(inputs)
         self.states = h, c
         if self.num_tasks == 1: