set __layout__ attribute for FusedRNNCell output states (#6747)

* set __layout__ attr for FucedRNNCell output states * right usage not that in desc
apache · Jun 20, 2017 · 94d8e18 · 94d8e18
1 parent 63ac793
commit 94d8e18
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/mxnet/rnn/rnn_cell.py b/python/mxnet/rnn/rnn_cell.py
@@ -672,11 +672,15 @@ def unroll(self, length, inputs, begin_state=None, layout='NTC', merge_outputs=N
                          mode=self._mode, name=self._prefix+'rnn',
                          **states)
 
+        attr = {'__layout__' : 'LNC'}
         if not self._get_next_state:
             outputs, states = rnn, []
         elif self._mode == 'lstm':
+            rnn[1]._set_attr(**attr)
+            rnn[2]._set_attr(**attr)
             outputs, states = rnn[0], [rnn[1], rnn[2]]
         else:
+            rnn[1]._set_attr(**attr)
             outputs, states = rnn[0], [rnn[1]]
 
         if axis == 1: