from tensorflow.contrib.learn.python.learn.graph_actions import infer

      random_seed.set_random_seed(self._config.tf_random_seed)
      contrib_framework.create_global_step(g)
      features, _ = input_fn()
      feed_dict = feed_fn() if feed_fn is not None else None
      predictions = self._get_predict_ops(features)
      if not isinstance(predictions, dict):
        predictions = {'predictions': predictions}
      # TODO(ipolosukhin): Support batching
      return infer(checkpoint_path, predictions, feed_dict=feed_dict)


class Estimator(BaseEstimator):
  """Estimator class is the basic TensorFlow model trainer/evaluator.

  Parameters:
    model_fn: Model function, takes features and targets tensors or dicts of
