Internal change

PiperOrigin-RevId: 296561671

official/nlp/optimization.py

@@ -137,16 +137,10 @@ class AdamWeightDecay(tf.keras.optimizers.Adam):
           use_locking=self._use_locking)
     return tf.no_op()
 
-  def apply_gradients(self,
-                      grads_and_vars,
-                      name=None,
-                      all_reduce_sum_gradients=True):
+  def apply_gradients(self, grads_and_vars, name=None):
     grads, tvars = list(zip(*grads_and_vars))
     (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)
-    return super(AdamWeightDecay, self).apply_gradients(
-        zip(grads, tvars),
-        name=name,
-        all_reduce_sum_gradients=all_reduce_sum_gradients)
+    return super(AdamWeightDecay, self).apply_gradients(zip(grads, tvars))
 
   def _get_lr(self, var_device, var_dtype, apply_state):
     """Retrieves the learning rate with the given state."""