Hi
I am trying to visualize intermediate layer outputs generate by one input image during inference of a PyTorch model. Preferably I would like to do this from a traced graph, for example one from the torchvision modelzoo.
I.e. if I have a model file created like this
import torch
import torchvision

org_model = torchvision.models.resnet18(pretrained=True)
traced_net = torch.jit.trace(org_model, torch.rand(1, 3, 224, 224))

torch.jit.save(traced_net, "resnet.pth") 

Then I want to be able to load that model and output the activations of for example “layer1”
traced_model_loaded = torch.jit.load("resnet.pth")

input_ = torch.rand(1, 3, 224, 224)
layer1_act = traced_model_loaded.layer1(input_)

Is this possible? If not, can I in some way modify the original PyTorch model so that an arbitrary number of layer activations becomes accessible?  Using forward hooks does not seem to be supported.
Thanks!