[MXNET-969] Fix buffer overflow in RNNOp (#12603)

* [MXNET-969] Fix buffer overflow in RNNOp Co-authored-by: Sina Md <[email protected]> * [MXNET-969] Use ternary op for statecell
apache · Sep 21, 2018 · 9f5c96e · 9f5c96e
1 parent cb6815c
commit 9f5c96e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/operator/rnn-inl.h b/src/operator/rnn-inl.h
@@ -576,7 +576,8 @@ class RNNOp : public Operator{
                        req[rnn_enum::kData],
                        req[rnn_enum::kParams],
                        req[rnn_enum::kState],
-                       req[rnn_enum::kStateCell],
+                       // State cell should be present for LSTMs, but is absent for other RNNs.
+                       param_.mode == rnn_enum::kLstm ? req[rnn_enum::kStateCell] : kNullOp,
                        param_.p,
                        param_.mode);
   }