Hi @Mona_Jalal
If you’d like to extract input and/or output from some modules in your model, you can easily achieve it by using ForwardHookManager in torchdistill 1.
Here is an example notebook  1 with ResNet.
In Section 3 of the notebook, you can use avgpool as the module path for torchdistill like forward_hook_manager.add_hook(model, 'avgpool', requires_input=False, requires_output=True)