Solved by techkang in post #6 
              
              
                I get your point. You can achieve you purpose by register_forward_hook, here is example code: 
from functools import partial

import torch as t
from torchvision.models import resnet50

model = resnet50()
data = t.randn(1, 3, 224, 224)

layer_output = {}


def get_all_layers(net):
    def hook_fn(m, …