export.py

import tensorflow.compat.v1 as tf

def export_model(estimator, export_dir, params,
                 checkpoint_path=None):


    def serving_input_receiver_fn():
        t = tf.placeholder(dtype=tf.int64,
                            shape=[1, params["n_ctx"]],
                            name='input_example_tensor')
        return tf.estimator.export.ServingInputReceiver(t, t)

    return estimator.export_saved_model(
        export_dir, serving_input_receiver_fn, checkpoint_path=checkpoint_path)