I am setting up FLOPs counter for models using register_forward_hook to register a counter at each module. However, I found forward hooks can only be applied on modules defined in the __init__ function of the model after some coding experiments, meaning we cannot register hooks on APIs like nn.functional.interpolate because they need to confirm the size of output which can be various if the input size is not fixed.
A post from Calculating flops of a given pytorch model 8 also mentioned this problem.
Are there any workarounds for registering a forward hook for those functions? Or are there any workarounds for counting FLOPs of interpolate together with other modules?