Keras like summary of a model

Weird, I may need to check my editor. It does not show up the reference of unet, but when I go to unet.py I did saw model_sizes is used…

It only works when I wrap a nn.Module inside a Sequential object, but what I want is to get output of every layer inside of a block. I am thinking how to hook at a m.children() level, takes BasicBlock as example it should have hook on all 7 layers inside it.