How can I obtain the output of an intermediate layer (feature extraction in tensor2tensor)?

Question

I need to see the intermediate values after infer, something similar to what is suggested [here][1] for Keras

https://keras.io/getting_started/faq/#how-can-i-obtain-the-output-of-an-intermediate-layer-feature-extraction

 model = ...  # create the original model

 layer_name = 'my_layer'
 intermediate_layer_model = keras.Model(inputs=model.input,
                                       outputs=model.get_layer(layer_name).output)
 intermediate_output = intermediate_layer_model(data) 

I cannot find a similar solution in tensor2tensor. Is there a hparam that I can set to cut my model at specific layer?