da4f2a9e6b7ff3045de54c8b470fa02f9aef169b,sonnet/src/momentum.py,Momentum,apply,#Momentum#,61

Before Change

          sparse_momentum_update = (mu * momentum.gather_nd(indices)) + update
          momentum.scatter_nd_update(indices, sparse_momentum_update)
          if self.use_nesterov:
            parameter.scatter_nd_sub(
                indices, (lr * update) + (lr * mu * sparse_momentum_update))
          else:
            parameter.scatter_nd_sub(indices, lr * sparse_momentum_update)
        else:
          momentum.assign((mu * momentum) + update)

After Change

          momentum.scatter_update(
              tf.IndexedSlices(sparse_momentum_update, indices))
          if self.use_nesterov:
            parameter.scatter_sub(tf.IndexedSlices(
                (lr * update) + (lr * mu * sparse_momentum_update), indices))
          else:
            parameter.scatter_sub(tf.IndexedSlices(
                lr * sparse_momentum_update, indices))
        else:

Instances