tf.estimator.predict slow with tensorflow ranking module

I'm currently using the TensorFlow Ranking module for a recommendation task. Specifically I'm modifying this tutorial file for my own purpose. I can't say the tutorial was very friendly to new users. As this is the first time I've interacted with TensorFlow, I'm just trying to make it run.

As you may notice, the example file doesn't say how to make predictions, so after I finished training the model, I modified its train_and_eval() function to predict. Here's my code.

def _train_op_fn(loss):
        """Defines train op used in ranking head."""
        update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
        minimize_op = optimizer.minimize(
                loss = loss, global_step = tf.compat.v1.train.get_global_step())
        train_op = tf.group([minimize_op, update_ops])
        return train_op

ranking_head = tfr.head.create_ranking_head(
        loss_fn = tfr.losses.make_loss_fn(loss),
        eval_metric_fns = get_eval_metric_fns(),
        train_op_fn = _train_op_fn
)

estimator =tf.estimator.Estimator(
        model_fn=tfr.model.make_groupwise_ranking_fn(
                group_score_fn=make_score_fn(),
                group_size = group_size,
                transform_fn=make_transform_fn(),
                ranking_head=ranking_head),
        config=tf.estimator.RunConfig(
                output_dir, save_checkpoints_steps=1000))

def predict_(feature_dict = {}):
    if feature_dict == {}:
        feature_dict = input_fn()
    pred_fn, pred_hook = get_eval_inputs(feature_dict)
    generator_ = estimator.predict(input_fn = pred_fn, hooks = [pred_hook])
    pred_list = list(generator_)

    return pred_list

The predict_ function takes dictionary

{'feature 1':[doc 1 score, doc 2 score...], 
 'feature 2':[doc 1 score, doc 2 score...], 
 ...}

and returns a list a scores for all docs in that order. (or at least that's what I think it should do)

The prediction results are very good. Problem is, it's really slow. Prediction for 400 docs takes more than 1 second (I only have 4 features). Is this a normal speed or there's space for optimization in the code? I heard tf.estimator reloads the graph everytime it makes a prediction, but I can't tell if it's the problem here.

from tf.estimator.predict slow with tensorflow ranking module