fix lstm memory shape

sony · May 24, 2019 · 1e5fae1 · 1e5fae1
1 parent 3cb68ed
commit 1e5fae1
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/src/nbla/cuda/cudnn/function/generic/gru.cu b/src/nbla/cuda/cudnn/function/generic/gru.cu
@@ -463,7 +463,7 @@ void GRUCudaCudnn<T>::setup_impl(const Variables &inputs,
   }
 
   // Set output shapes
-  outputs[0]->reshape({seq_len_, batch_size, hidden_size_, num_directions_},
+  outputs[0]->reshape({seq_len_, batch_size, num_directions_ * hidden_size_},
                       true);
   outputs[1]->reshape(inputs[1]->shape(), true);
 }

diff --git a/src/nbla/cuda/cudnn/function/generic/lstm.cu b/src/nbla/cuda/cudnn/function/generic/lstm.cu
@@ -450,7 +450,7 @@ void LSTMCudaCudnn<T>::setup_impl(const Variables &inputs,
   }
 
   // Set output shapes
-  outputs[0]->reshape({seq_len_, batch_size, hidden_size_, num_directions_},
+  outputs[0]->reshape({seq_len_, batch_size, num_directions_ * hidden_size_},
                       true);
   outputs[1]->reshape(inputs[1]->shape(), true);
   outputs[2]->reshape(inputs[2]->shape(), true);

diff --git a/src/nbla/cuda/cudnn/function/generic/rnn.cu b/src/nbla/cuda/cudnn/function/generic/rnn.cu
@@ -433,7 +433,7 @@ void RNNCudaCudnn<T>::setup_impl(const Variables &inputs,
   }
 
   // Set output shapes
-  outputs[0]->reshape({seq_len_, batch_size, hidden_size_, num_directions_},
+  outputs[0]->reshape({seq_len_, batch_size, num_directions_ * hidden_size_},
                       true);
   outputs[1]->reshape(inputs[1]->shape(), true);
 }