Hello everyone,
I would like to prune my model and then script it but I receive the following error: AttributeError: 'LnStructured' object has no attribute '__name__'
Can we prune a model and then script it?
I see that LnStructured actually add a forward pre_hook but jit.script can’t resolve its name.
cf:
Traceback (most recent call last):
  File "src/model_optimizer.py", line 226, in <module>
    script_module(pruned_model)
  File "src/model_optimizer.py", line 20, in script_module
    module = jit.script(module)
  File "/venv/lib/python3.8/site-packages/torch/jit/_script.py", line 1257, in script
    return torch.jit._recursive.create_script_module(
  File "/venv/lib/python3.8/site-packages/torch/jit/_recursive.py", line 451, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/venv/lib/python3.8/site-packages/torch/jit/_recursive.py", line 513, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/venv/lib/python3.8/site-packages/torch/jit/_script.py", line 587, in _construct
    init_fn(script_module)
  File "/venv/lib/python3.8/site-packages/torch/jit/_recursive.py", line 491, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/venv/lib/python3.8/site-packages/torch/jit/_recursive.py", line 513, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/venv/lib/python3.8/site-packages/torch/jit/_script.py", line 587, in _construct
    init_fn(script_module)
  File "/venv/lib/python3.8/site-packages/torch/jit/_recursive.py", line 491, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/venv/lib/python3.8/site-packages/torch/jit/_recursive.py", line 465, in create_script_module_impl
    hook_stubs, pre_hook_stubs = get_hook_stubs(nn_module)
  File "/venv/lib/python3.8/site-packages/torch/jit/_recursive.py", line 758, in get_hook_stubs
    if pre_hook.__name__ in hook_map:
AttributeError: 'LnStructured' object has no attribute '__name__'

Here is the code to reproduce the error :
import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
import torch.jit as jit

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.layers = nn.Sequential(
            nn.Linear(10, 10),
            nn.ReLU(),

            nn.Linear(10, 10),
            nn.Sigmoid(),
        )

    def forward(self, t):
        return self.layers(t)

model = Model()

for _, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        prune.ln_structured(module, name='weight', amount=0.4, dim=1, n=float('-inf'))

jit.script(model)

Thanks,
Thytu