Update rnn_cell.py

oza · Apr 4, 2017 · 5c69214 · 5c69214
1 parent 911f747
commit 5c69214
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/python/mxnet/rnn/rnn_cell.py b/python/mxnet/rnn/rnn_cell.py
@@ -736,14 +736,16 @@ def unroll(self, length, inputs, begin_state=None, layout='NTC', merge_outputs=N
             begin_state = self.begin_state()
 
         p = 0
+        next_states = []
         for i, cell in enumerate(self._cells):
             n = len(cell.state_shape)
             states = begin_state[p:p+n]
             p += n
             inputs, states = cell.unroll(length, inputs=inputs, begin_state=states, layout=layout,
                                          merge_outputs=None if i < num_cells-1 else merge_outputs)
+            next_states.extend(states)
 
-        return inputs, states
+        return inputs, next_states
 
 
 class DropoutCell(BaseRNNCell):