Construct the data, model, loss and optimizer then train.

def build_and_train(iterations, log_stride, test=False):
  """Construct the data, model, loss and optimizer then train."""

  # Test mode settings.
  batch_size = 2 if test else FLAGS.batch_size
  num_mems = 2 if test else FLAGS.num_mems
  num_heads = 1 if test else FLAGS.num_mems
  num_blocks = 1 if test else FLAGS.num_mems
  head_size = 4 if test else FLAGS.head_size
  max_length = 3 if test else FLAGS.max_length
  max_nest = 2 if test else FLAGS.max_nest
  mlp_size = (20,) if test else (256, 256, 256, 256)

  with tf.Graph().as_default():
    t0 = time.time()

    # Initialize the dataset.
    lte_train = learn_to_execute.LearnToExecute(
        batch_size, max_length, max_nest)
    lte_test = learn_to_execute.LearnToExecute(
        batch_size, max_length, max_nest, mode=learn_to_execute.Mode.TEST)
    train_data_iter = lte_train.make_one_shot_iterator().get_next()
    test_data_iter = lte_test.make_one_shot_iterator().get_next()
    output_size = lte_train.state.vocab_size

    # Create the model.
    core = snt.RelationalMemory(
        mem_slots=num_mems,
        head_size=head_size,
        num_heads=num_heads,
        num_blocks=num_blocks,
        gate_style=FLAGS.gate_style)
    final_mlp = snt.nets.MLP(
        output_sizes=mlp_size,
        activate_final=True)
    model = SequenceModel(
        core=core,
        target_size=output_size,
        final_mlp=final_mlp)
    tf.logging.info("Instantiated models ({:3f})".format(time.time() - t0))

    # Define the loss & accuracy.
    def loss_fn(inputs, targets, input_sequence_length, output_sequence_length):
      """Creates the loss and the exports."""
      logits = model(
          inputs, targets, input_sequence_length, output_sequence_length)
      targets = tf.cast(targets, tf.int32)
      sq_sz_out_max = targets.shape[0].value

      # Create a mask to ignore accuracy on buffer characters.
      sequence_sizes = tf.cast(output_sequence_length, tf.float32)
      lengths_transposed = tf.expand_dims(sequence_sizes, 1)
      range_row = tf.expand_dims(
          tf.range(0, sq_sz_out_max, 1, dtype=tf.float32), 0)
      mask = tf.cast(tf.transpose(tf.less(range_row, lengths_transposed)),
                     tf.float32)

      # Compute token accuracy and solved.
      correct = tf.equal(tf.argmax(logits, 2), tf.argmax(targets, 2))
      solved = tf.reduce_all(tf.boolean_mask(correct, tf.squeeze(mask)), axis=0)
      token_acc = tf.reduce_sum(tf.cast(correct, tf.float32) * mask)
      token_acc /= tf.reduce_sum(sequence_sizes)

      # Compute Loss.
      mask = tf.cast(tf.tile(tf.expand_dims(mask, 2), (1, 1, logits.shape[2])),
                     tf.float32)
      masked_logits = logits * mask
      masked_target = tf.cast(targets, tf.float32) * mask
      logits_flat = tf.reshape(masked_logits,
                               [sq_sz_out_max * batch_size, -1])
      target_flat = tf.reshape(masked_target,
                               [sq_sz_out_max * batch_size, -1])
      xent = tf.nn.softmax_cross_entropy_with_logits(logits=logits_flat,
                                                     labels=target_flat)
      loss = tf.reduce_mean(xent)
      return loss, token_acc, solved

    # Get training step counter.
    global_step = tf.train.get_or_create_global_step()

    # Create the optimizer.
    learning_rate_op = tf.reduce_max([
        tf.train.exponential_decay(
            FLAGS.learning_rate,
            global_step,
            decay_steps=FLAGS.epochs // 100,
            decay_rate=0.9,
            staircase=False),
        FLAGS.min_learning_rate
    ])
    optimizer = tf.train.AdamOptimizer(learning_rate_op)

    # Compute loss, accuracy & the step op.
    inputs, targets, _, input_lengths, output_lengths = train_data_iter
    train_loss, train_acc, train_sol = loss_fn(
        inputs, targets, input_lengths, output_lengths)
    step_op = optimizer.minimize(train_loss, global_step=global_step)
    inputs, targets, _, input_lengths, output_lengths = test_data_iter
    _, test_acc, test_sol = loss_fn(
        inputs, targets, input_lengths, output_lengths)
    tf.logging.info("Created losses and optimizers ({:3f})".format(
        time.time() - t0))

    # Begin Training.
    t0 = time.time()
    tf.logging.info("Starting training ({:3f})".format(time.time() - t0))
    with tf.train.SingularMonitoredSession() as sess:
      for it in six.moves.range(iterations):
        sess.run([step_op, learning_rate_op])
        if it % log_stride == 0:
          loss_v, train_acc_v, test_acc_v, train_sol_v, test_sol_v = sess.run([
              train_loss, train_acc, test_acc, train_sol, test_sol])
          elapsed = time.time() - t0
          tf.logging.info(
              "iter: {:2d}, train loss {:3f}; train acc {:3f}; test acc {:3f};"
              " train solved {:3f}; test solved {:3f}; ({:3f})".format(
                  it, loss_v, train_acc_v, test_acc_v, train_sol_v, test_sol_v,
                  elapsed))