Output intermediate layer from the torch.jit.ScriptModule version of the network

Hi guys,

I have the following problem. I use Torch Script for saving the trained model. Afterwards I need to load it and
take the gradient of the model output with respect to one of the intermediate layers output. At the same time I don’t want to output this intermediate layer with the forward method. As far as I understand register_forward_hook method is not yet implemented in ScriptModule. Could you please tell, if there is another way to solve this problem?

Best,
Andrei