Me gustaría gestionar mi entrenamiento con untf.estimator.Estimator pero tiene algunos problemas para usarlo junto con eltf.data API

Tengo algo como esto:

def model_fn(features, labels, params, mode):
  # Defines model's ops.
  # Initializes with tf.train.Scaffold.
  # Returns an tf.estimator.EstimatorSpec.

def input_fn():
  dataset = tf.data.TextLineDataset("test.txt")
  # map, shuffle, padded_batch, etc.

  iterator = dataset.make_initializable_iterator()

  return iterator.get_next()

estimator = tf.estimator.Estimator(model_fn)
estimator.train(input_fn)

Como no puedo usar unmake_one_shot_iterator para mi caso de uso, mi problema es queinput_fn contiene un iterador que debe inicializarse dentro demodel_fn (aquí, yo usotf.train.Scaffold para inicializar operaciones locales).

Además, entendí que no solo podemos usarinput_fn = iterator.get_next de lo contrario, las otras operaciones no se agregarán al mismo gráfico.

¿Cuál es la forma recomendada de inicializar el iterador?