back-propagate error to the second head of nerv.MSELayer

1 files changed, 3 insertions, 0 deletions

diff --git a/nerv/layer/mse.lua b/nerv/layer/mse.lua
index c1ea596..594eb0c 100644
--- a/nerv/layer/mse.lua
+++ b/nerv/layer/mse.lua
@@ -63,9 +63,12 @@ function MSELayer:back_propagate(bp_err, next_bp_err, input, output, t)
         t = 1
     end
     local nbe = next_bp_err[1]
+    local nbe2 = next_bp_err[2]
     nbe:add(nbe, self.diff[t], 0.0, self.scale)
+    nbe2:add(nbe2, self.diff[t], 0.0, -self.scale)
     if bp_err[1] ~= nil then
         nbe:scale_rows_by_col(bp_err[1])
+        nbe2:scale_rows_by_col(bp_err[1])
     end
 end