Is there a way to speed up the call or prediction function for a Keras model while monitoring the tensors?

Question

I\'m trying to make predictions with a model I built in keras, but I need it to output tensors because I\'m using reinforcement learning and so I need to keep track of the g