Visual trace for nn.Module forward call

Hi everyone,

Do you know of any visualizers/debuggers that display the visual trace of all forward calls in a module along with their results (e.g. result tensor shape for each forward call)? Consider a complex instance of nn.Module like BertModel. I’d love to pass input to this module and get a visual trace (a simple printout would be a great start) of all forward calls executed and their corresponding tensor shapes.

I can imagine someone already made a similar debugger using hooks - i.e. using register_forward_hook recursively attached to every layer in the module.