import tensorflow as tf

    train_op = self.optimize(loss, num_async_replicas=num_async_replicas)

    if common_layers.is_on_tpu():
      _remove_summaries()  # summaries not currently working on TPU
      return tf.contrib.tpu.TPUEstimatorSpec(
          tf.estimator.ModeKeys.TRAIN, loss=loss, train_op=train_op)
    else:
      return tf.estimator.EstimatorSpec(
