import tensorflow as tf

                    mask_loss_t, learning_rate_t])
  else:
    train_op = None
    scaffold_fn = None

  return tf.contrib.tpu.TPUEstimatorSpec(
      mode=mode,
      loss=total_loss,
      train_op=train_op,
      host_call=host_call,
