🚀 Feature

Hello,

This feature is related to a thread from discuss forum

I would like to propose mechanism of saving _forward_pre_hooks and _forward_hooks while model is preparing for quantization.

Motivation

There are cases that depend on such functionality. Here I duplicate what I said about it in the forum

Have a look at this example of pre forward hook run. Here’s EfficientNet implementation that assumes it can be integrated as backbone to FPN (https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/efficientnet.py#L471)
So here we can see that during forward blocks are called sequentially than we collect input features from specific layers.

Pitch

We can add several lines here https://github.com/pytorch/pytorch/blob/master/torch/quantization/quantize.py#L337
to preserve mod's hooks in new_mod

Additional context

Additionally it is needed to discuss the possibilities of preserving hooks while fusing modules

cc @jerryzh168 @jianyuh @dzhulgakov @raghuramank100 @jamesr66a