Internal change

PiperOrigin-RevId: 296550109

official/nlp/optimization.py

@@ -137,10 +137,16 @@ class AdamWeightDecay(tf.keras.optimizers.Adam):
           use_locking=self._use_locking)
     return tf.no_op()
 
-  def apply_gradients(self, grads_and_vars, name=None):
+  def apply_gradients(self,
+                      grads_and_vars,
+                      name=None,
+                      all_reduce_sum_gradients=True):
     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))
+    return super(AdamWeightDecay, self).apply_gradients(
+        zip(grads, tvars),
+        name=name,
+        all_reduce_sum_gradients=all_reduce_sum_gradients)
 
   def _get_lr(self, var_device, var_dtype, apply_state):
     """Retrieves the learning rate with the given state."""