Keeping the [1, 3, 224, 224] input beside the capture is a useful boundary. For a hierarchical model, I would add one nearby non-square or non-window-aligned input and compare only the stage-boundary shapes first. Padding and window partition behavior can remain hidden in an otherwise familiar module overview.
The same breadcrumbs would then let a reader locate the first boundary whose shape differs before opening hundreds of operations. Labeling each capture with input shape and evaluation mode would keep that comparison concrete without implying that registration order or one successful trace describes every execution.