how to restore a meta graph with custom ops for inference in tensorflow

[karterotte]

I used a custom op for data transform and got a saved model. All above is ok Now I want to load this model and do some predict. Code likes (ps_pull):

with tf.Session(config=config) as sess:
          # init
          self.do_prepare_data()
          
          # load pb
          meta_graph_def = tf.saved_model.loader.load(
              sess, ["serve"], model_path)
          signature = meta_graph_def.signature_def
          ...
          # load ps   this is my custom op 
          ps_pull = tf.load_op_library('/usr/local/src/ps_lite.so')
          sess.run(ps_pull)
          
          # load data 
          train_input, sample_id, _ = dataset_to_input(self.train_data)
         
          # here the data contains custom op
          sess.run(train_input)

But the data is not correct.It seems the custom op didn't work and return init values.

How do I fix it? Tanks