Failed to profile torchscript models

[terryzhao127]

I tried to apply this nice profiler to torchscript models but it failed:

Traceback (most recent call last):
  File "/home/guikarist/projects/test-torch/profiler/test_flops_profile.py", line 48, in <module>
    flops, macs, params = get_model_profile(
  File "/home/guikarist/.miniconda3/envs/pytorch-2.0.1/lib/python3.9/site-packages/flops_profiler/profiler.py", line 1505, in get_model_profile
    prof.start_profile(ignore_list=ignore_modules)
  File "/home/guikarist/.miniconda3/envs/pytorch-2.0.1/lib/python3.9/site-packages/flops_profiler/profiler.py", line 152, in start_profile
    self.model.apply(
  File "/home/guikarist/.miniconda3/envs/pytorch-2.0.1/lib/python3.9/site-packages/torch/nn/modules/module.py", line 884, in apply
    module.apply(fn)
  File "/home/guikarist/.miniconda3/envs/pytorch-2.0.1/lib/python3.9/site-packages/torch/nn/modules/module.py", line 885, in apply
    fn(self)
  File "/home/guikarist/.miniconda3/envs/pytorch-2.0.1/lib/python3.9/site-packages/flops_profiler/profiler.py", line 104, in register_module_hooks
    module.__pre_hook_handle__ = module.register_forward_pre_hook(
  File "/home/guikarist/.miniconda3/envs/pytorch-2.0.1/lib/python3.9/site-packages/torch/jit/_script.py", line 942, in fail
    raise RuntimeError(name + " is not supported on ScriptModules")
RuntimeError: register_forward_pre_hook is not supported on ScriptModules

Is there anyway to solve the problem? This issue may be related.

[terryzhao127]

FYI, register_forward_pre_hook post-scripted models has not been supported yet: pytorch/pytorch#29037 (comment)