Update for those wanting to do something similar to me.
I found kind of a workaround. I registered a forward hook to all “leaf” modules in the model, that is modules that do not have submodules. The hook/function creates a storage object with information (module, in data, out data), this object is put into a global list.
The list if updated in order of the forward pass. Once a forward pass is complete, I have a list of storage objects in order of the forward pass with all the info I need.
There are a few clear weaknesses however, this method can only register hooks on modules that have been initialized. If the model uses i.e. torch.nn.functional.max_pool2d or any other method from torch.nn.functional in the forward method of the model, they won’t show up in model.modules() and won’t get a forward hook.
Secondly, the global list may cause several problems: Unsure how it will work in a parallel setup (multiple GPUs) or even just networks that have branching such as skip connections.