Модель Tensorflow выводит другое значение во время вывода

Во время обучения я записываю прогнозируемые значения, которые моя модель регрессора выводит на основе данных обучения. Когда я запустил тот же набор данных в режиме прогнозирования, диапазон значений, которые выводит модель, сильно отличается:

Тензорборд

Здесь мы видим, что модель предсказывает значения в диапазоне (140, 250).

Прогнозы по тому же набору данных

Здесь у нас есть модель, предсказывающая значения между (17, 23). Что дает?

Я подозреваю, что API-интерфейс оценщика волшебным образом не сохраняет moving_mean и moving_variance при использовании tf.layers.batch_normalization.

Моя модель:

def model_fn(features, labels, mode, params):
  training = mode == tf.estimator.ModeKeys.TRAIN
  extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

  x = tf.reshape(features, [-1, 32, 32, 3])
  x = tf.layers.batch_normalization(x, training=training, name='norm_128')

  i = 1
  for filters in [32, 64]:
    x = tf.layers.conv2d(x, filters=filters, kernel_size=3, activation=None, name='conv{}'.format(i))
    x = tf.layers.batch_normalization(x, training=training, name='norm{}'.format(i))
    x = tf.nn.relu(x, name='act{}'.format(i))
    i += 1

    x = tf.layers.conv2d(x, filters=filters * 2, kernel_size=3, strides=2, activation=None, name='pool{}'.format(i))
    x = tf.layers.batch_normalization(x, training=training, name='norm{}'.format(i))
    x = tf.nn.relu(x, name='act{}'.format(i))
    i += 1

  flat = tf.contrib.layers.flatten(x, scope='flatten')
  dropout = tf.layers.dropout(flat, rate=params['dropout_rate'], training=training, name='dropout')
  output_layer = tf.layers.dense(dropout, units=1, name='output_layer')

  predictions = tf.reshape(output_layer, [-1])

  predictions_dict = {
    'pred': predictions,
  }

  if mode == tf.estimator.ModeKeys.PREDICT:
    return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions_dict)

  loss = tf.losses.mean_squared_error(labels=labels, predictions=predictions)

  tf.summary.scalar('loss', loss)
  tf.summary.histogram('prediction', predictions)
  tf.summary.scalar('prediction', tf.reduce_mean(predictions))

  optimizer = tf.train.AdamOptimizer(learning_rate=params['learning_rate'])
  with tf.control_dependencies(extra_update_ops):
    train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
  eval_metric_ops = {
    'rmse_val': tf.metrics.root_mean_squared_error(labels=tf.cast(labels, tf.float32), predictions=predictions)
  }

  tf.summary.scalar('rmse_train', eval_metric_ops['rmse_val'][1])


  return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op, eval_metric_ops=eval_metric_ops)