fix CuDNN RecurrentOp Gradient init
Summary: CuDNN RecurrentNet GradientOp did not pass the DROPOUT information to the initializer, causing incorrect scratch space size to be estimated. We have an assertion encorcing that scratch space is same for forward and backward ops, so this failed an assertion. We currently hard-code dropout to be 1.0, so this has had no effect on correctness in our tests. For some reason with num_layers=1 there wasn't an issue, but with num_layers>=2, the scratch space size was different.
Reviewed By: salexspb
Differential Revision: D4904715
fbshipit-source-id: 780266c5ecf1f7a32387edcb6fc498a13ac782ac
diff --git a/caffe2/operators/recurrent_op_cudnn.cc b/caffe2/operators/recurrent_op_cudnn.cc
index 6951852..e36c2e3 100644
--- a/caffe2/operators/recurrent_op_cudnn.cc
+++ b/caffe2/operators/recurrent_op_cudnn.cc
@@ -295,7 +295,7 @@
bool RecurrentGradientOp<T>::RunOnDevice() {
const int seqLength = Input(INPUT).dim32(0);
if (Input(INPUT).dims() != cachedInputDims_) {
- initialize(Input(INPUT));
+ initialize(Input(INPUT), Output(DROPOUT_STATES));
cachedInputDims_ = Input(INPUT).dims();
}
CUDNN_ENFORCE(cudnnGetRNNTrainingReserveSize(
