How to obtain the intermediate features (bottleneck features) in the downsampling stage of U-Net?

[Guoxt]

intermediate_outputs = {}
def forward_hook(layer_name):
def hook(module, input, output):
intermediate_outputs[layer_name] = output

return hook

input = torch.randn(4,1,64,64,64)
net.model[?].register_forward_hook(forward_hook("XXX"))

How should this be set?