[layer] LSTMCell bug fix

This patch applies bug fix for lstmcell.

Signed-off-by: Parichay Kapoor <pk.kapoor@samsung.com>

diff --git a/nntrainer/layers/lstmcell.cpp b/nntrainer/layers/lstmcell.cpp
index 3e01cb89202e3ce2f7f7a56e7570db00d0bcc950..24927ffe8c4ce400a2a4f7fcdcc317b30a43d99b 100644 (file)
--- a/nntrainer/layers/lstmcell.cpp
+++ b/nntrainer/layers/lstmcell.cpp
@@ -354,6 +354,8 @@ void LSTMCellLayer::calcGradient(RunLayerContext &context) {
     dc.multiply_strided(hf, dc_nx);
     Tensor cs_prev = m_cell_.getBatchSlice(start_timestep - 1, 1);
     dc.multiply_strided(cs_prev, dhf);
+  } else {
+    dhf.setZero();
   }
 
   dc.multiply_strided(hg, dhi);

diff --git a/nntrainer/tensor/tensor.cpp b/nntrainer/tensor/tensor.cpp
index c868c3de808a4f69b0a57b9d11657897637ba563..f096276d695adf31b19b59b590d6fa757dd88951 100644 (file)
--- a/nntrainer/tensor/tensor.cpp
+++ b/nntrainer/tensor/tensor.cpp
@@ -1367,7 +1367,12 @@ void Tensor::setValue(float val) {
   std::fill(data, data + size(), val);
 }
 
-void Tensor::setZero() { sscal(size(), 0, getData(), 1); }
+void Tensor::setZero() {
+  if (contiguous)
+    sscal(size(), 0, getData(), 1);
+  else
+    apply_i([](float val) -> float { return 0; });
+}
 
 std::vector<unsigned int> Tensor::argmax() const {
   const float *data = getData();