Getting out shape after each module in model_summary

I was messing around in the 11_train_imagenette notebook, mostly with model_summary to see what all of the layers are doing, but the ResBlocks are being difficult. I added:

def is_res_layer(l):
    res_layers = (ResBlock)
    return isinstance(l,res_layers)

and modified model_summary a little to:

def model_summary(model, find=None):
    xb,yb = get_batch(data.valid_dl, learn)
    mods = find_modules(model, find) if find else model.children()
    f = lambda hook,mod,inp,out: print(f'in: {inp[0].shape}\n{mod}\nout: {out.shape}\n{"-"*90}')
    with Hooks(mods, f) as hooks: learn.model(xb)

so now model_summary(learn.model) shows this:


which i like until it gets to the next Sequential with the ResBlocks and the out shape doesn’t show again until it goes through that whole Sequential which has three ResBlocks.

What I’ve got so far is to do:
model_summary(learn.model, is_res_layer)


which shows the out shape after the ResBlock, but doesn’t show the out shape after the Sequentials in the ResBlock. It also only shows the ResBlocks in the model.

I want to keep the format like what learn.model prints out, but am stuck on how to get the out shape to show after each Sequential, ResBlock, and Sequential in a ResBlock. If you have any ideas lmk!