Predict from loaded BERT model

I was trying to make a prediction from a loaded tensorflow
model. Though I’m not sure if it’s correct how I previously saved
it, specifically I have doubts about code inside serving_input_fn()
function (MAX_SEQ_LENGTH=128):

def serving_input_fn(): feature_spec = { "input_ids" : tf.FixedLenFeature([MAX_SEQ_LENGTH], tf.int64), "input_mask" : tf.FixedLenFeature([MAX_SEQ_LENGTH], tf.int64), "segment_ids" : tf.FixedLenFeature([MAX_SEQ_LENGTH], tf.int64), "label_ids" : tf.FixedLenFeature([], tf.int64) } serialized_tf_example = tf.placeholder(dtype=tf.string,shape=[None],name='input_example_tensor') receiver_tensors = {'example': serialized_tf_example} features = tf.parse_example(serialized_tf_example, feature_spec) return tf.estimator.export.ServingInputReceiver(features, receiver_tensors) estimator.export_saved_model('gs://bucket/trained_model, serving_input_receiver_fn=serving_input_fn) 

When I try to predict from loaded model:

from tensorflow.contrib import predictor predict_fn = predictor.from_saved_model(LOAD_PATH) input_features_test = convert_examples_to_features( test_examples,label_list, MAX_SEQ_LENGTH, tokenizer) predictions = predict_fn({'example':input_features_test[0]}) 

it returns this error:

ValueError: Cannot feed value of shape () for Tensor
‘input_example_tensor:0’, which has shape ‘(?,)’

How should I change serving_input_fn() method?

If you want to reproduce it: github_repo (you
should download variables from
here and put it in trained_model/1608370941/ folder)

This is the tutorial I followed to fine tune BERT model on
google cloud TPU.