Почему эта простая модель tf.keras не обучается после преобразования в оценщик тензорного потока?

Я пытаюсь преобразовать модель tf.keras в оценщик тензорного потока с помощью tf.keras.estimator.model_to_estimator, но полученный оценщик, похоже, не поддается обучению.

Я пробовал моделировать y = (x_1 + x_2) / 2, используя как последовательные, так и функциональные API-интерфейсы tf.keras, и хотя модели tf.keras работают отлично, ни одна из них не работает после преобразования в оценщики. Использование tf.estimator.LinearRegressor с теми же функциями ввода работает, поэтому я не думаю, что проблема связана с функциями ввода.

Вот минимальный рабочий пример для последовательно определенной модели tf.keras:

import numpy as np
import tensorflow as tf
import functools

sample_size = 1000

x_train = np.random.randn(sample_size, 2).astype(np.float32)
y_train = np.mean(x_train, axis=1).astype(np.float32) 

x_test = np.random.randn(sample_size, 2).astype(np.float32)
y_test = np.mean(x_test, axis=1).astype(np.float32) 

model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Dense(1, input_shape=(2,), name="Prediction"))
adam = tf.keras.optimizers.Adam(lr=0.1)
model.compile(loss='MSE', optimizer=adam)
#model.fit(x=x_train, y=y_train, epochs=10, batch_size=64)  # This works

est = tf.keras.estimator.model_to_estimator(keras_model=model)

def train_input_fn(batch_size):
    dataset = tf.data.Dataset.from_tensor_slices(({"Prediction_input": x_train}, y_train))
    return dataset.shuffle(sample_size).batch(batch_size).repeat()

def eval_input_fn(batch_size):
    dataset = tf.data.Dataset.from_tensor_slices(({"Prediction_input": x_test}, y_test))
    return dataset.batch(batch_size)

est.train(input_fn=functools.partial(train_input_fn, 64), steps=10)

eval_metrics = est.evaluate(input_fn=functools.partial(eval_input_fn, 1))
print('Evaluation metrics:', eval_metrics)

Оценщик обучен 10 шагам, которых должно быть более чем достаточно, чтобы снизить потери. Насколько я могу судить, увеличение количества шагов не имеет значения.

Когда я запускаю это в тензорном потоке 1.5.0, я получаю предупреждение о вызове reduce_mean, когда keep_dims устарел, когда модель tf.keras скомпилирована, но она отлично обучается как есть.

Это ошибка или я что-то упускаю?