892bdaa5bc8b93380e3f89cf716c6f2626f7f366,sonnet/src/sgd.py,SGD,apply,#SGD#Any#Any#,35

Before Change

      if update is not None:
        optimizer_utils.check_same_dtype(update, parameter)
        learning_rate = tf.cast(self.learning_rate, update.dtype.base_dtype)
        parameter.assign_sub(update * learning_rate)


class FastSGD(base.Module):
  Faster Stochastic Gradient Descent (SGD) module.

After Change

      if update is not None:
        optimizer_utils.check_same_dtype(update, parameter)
        learning_rate = tf.cast(self.learning_rate, update.dtype.base_dtype)
        if isinstance(update, tf.IndexedSlices):
          parameter.scatter_nd_sub(
              update.indices, update.values * learning_rate)
        else:
          parameter.assign_sub(update * learning_rate)


class FastSGD(base.Module):
  Faster Stochastic Gradient Descent (SGD) module.

  def __init__(self, learning_rate, name=None):

Instances