it’s much easier to write a piece of code
import torch
import torchvision

def hook_fn(m, i, o):
    print(o.shape)

model = torchvision.models.alexnet(pretrained=True)
model.classifier[4].register_forward_hook(hook_fn)

a = torch.zeros(1, 3, 256, 256)
model.eval()
with torch.no_grad():
    model(a)