def tff_model_builder():
  return tff.learning.from_keras_model(
      keras_model=keras_model_builder_with_zeros(),
      input_spec=get_input_spec(),
      loss=tf.keras.losses.MeanSquaredError(),
      metrics=metrics_builder())