add proper gradient summation to shift operator

1 files changed, 4 insertions, 2 deletions

diff --git a/src/graph/node_operators_unary.h b/src/graph/node_operators_unary.h
index 0fc17d28..2cc4fa37 100644
--- a/src/graph/node_operators_unary.h
+++ b/src/graph/node_operators_unary.h
@@ -997,11 +997,13 @@ struct ShiftNodeOp : public UnaryNodeOp {
       : UnaryNodeOp(a, a->shape()), shift_(shift) {}
 
   NodeOps forwardOps() {
-    return {NodeOp(Shift(val_, child(0)->val(), shift_, false))};
+    // last parameter beta=0 says to use = (out = in + beta * out)
+    return {NodeOp(Shift(val_, child(0)->val(), shift_, false, 0.f))};
   }
 
   NodeOps backwardOps() {
-    return {NodeOp(Shift(child(0)->grad(), adj_, shift_, true))};
+    // last parameter beta=1 says to use += (out = in + beta * out)
+    return {NodeOp(Shift(child(0)->grad(), adj_, shift_, true, 1.0f))};
   }
 
   const std::string type() { return "shift"; }