Internal change

PiperOrigin-RevId: 340748908

official/nlp/tasks/electra_task.py

@@ -215,7 +215,6 @@ class ElectraPretrainTask(base_task.Task):
           aux_losses=model.losses)
       # Scales loss as the default gradients allreduce performs sum inside the
       # optimizer.
-      # TODO(b/154564893): enable loss scaling.
       scaled_loss = loss / tf.distribute.get_strategy().num_replicas_in_sync
     tvars = model.trainable_variables
     grads = tape.gradient(scaled_loss, tvars)