import tensorflow as tf

            shape[1] = hparams.max_length
          v.set_shape(shape)
      else:
        shape = logits.get_shape().as_list()
        if shape[0] is None:
          shape[0] = params["batch_size"]
        if shape[1] is None:
          shape[1] = hparams.max_length
        logits.set_shape(shape)

    assert "training" in losses_dict

    # Summarize losses
    with tf.name_scope("losses"):
      for loss_name, loss_val in losses_dict.items():
        tf.summary.scalar(loss_name, loss_val)

    # Accumulate losses
    loss = sum(losses_dict.values())

    # EVAL mode
    if mode == tf.estimator.ModeKeys.EVAL:
      return model.estimator_spec_eval(features, logits, labels, loss,
                                       losses_dict)

    # TRAIN mode
    assert mode == tf.estimator.ModeKeys.TRAIN
    num_async_replicas = (1 if (use_tpu or not config) else
                          config.t2t_device_info["num_async_replicas"])
    return model.estimator_spec_train(
