Merged commit includes the following changes: (#7100)

254874613  by hongkuny<hongkuny@google.com>:

    Update glue tasks enum to match directory name

--
254866171  by taylorrobie<taylorrobie@google.com>:

    Internal change

PiperOrigin-RevId: 254874613

official/bert/create_finetuning_data.py

@@ -40,8 +40,8 @@ flags.DEFINE_string(
     "The input data dir. Should contain the .tsv files (or other data files) "
     "for the task.")
 
-flags.DEFINE_enum("classification_task_name", "mnli",
-                  ["cola", "mnli", "mrpc", "xnli"],
+flags.DEFINE_enum("classification_task_name", "MNLI",
+                  ["COLA", "MNLI", "MRPC", "XNLI"],
                   "The name of the task to train BERT classifier.")
 
 # BERT Squad task specific flags.

official/bert/optimization.py

@@ -136,14 +136,24 @@ class AdamWeightDecay(tf.keras.optimizers.Adam):
 
   def _resource_apply_dense(self, grad, var):
     var_dtype = var.dtype.base_dtype
+
+    try:
+      lr_t = self.apply_cache[var.device, var.dtype.base_dtype].lr_t
+    except AttributeError:
       lr_t = self._decayed_lr_t[var_dtype]
+
     with tf.control_dependencies([self._decay_weights_op(var, lr_t)]):
       return super(AdamWeightDecay, self)._resource_apply_dense(
           grad, var)
 
   def _resource_apply_sparse(self, grad, var, indices):
     var_dtype = var.dtype.base_dtype
+
+    try:
+      lr_t = self.apply_cache[var.device, var.dtype.base_dtype].lr_t
+    except AttributeError:
       lr_t = self._decayed_lr_t[var_dtype]
+
     with tf.control_dependencies([self._decay_weights_op(var, lr_t)]):
       return super(AdamWeightDecay, self)._resource_apply_sparse(
           grad, var, indices)