TIMM tresnet_l model fails with Python builtin <built-in method apply ... (InplaceABN.apply?)> is currently not supported in TorchScript

[xwang233]

🐛 Describe the bug

Reproduce:

root@c73318efaa9b:/workspace/timm-models/pytorch-image-models# python -u benchmark.py --bench train --model tresnet_l --img-size 224 -b 128 --torchscript --fuser nvfuser  
Benchmarking in float32 precision. NCHW layout. torchscript enabled
Model tresnet_l created, param count: 55989256
ERROR: "
Python builtin <built-in method apply of FunctionMeta object at 0x55c069d6eaa0> is currently not supported in Torchscript:
  File "/opt/conda/lib/python3.8/site-packages/inplace_abn/functions.py", line 153
def inplace_abn(x, weight, bias, running_mean, running_var,
                training=True, momentum=0.1, eps=1e-05, activation="leaky_relu", activation_param=0.01):
    return InPlaceABN.apply(x, weight, bias, running_mean, running_var,
           ~~~~~~~~~~~~~~~~ <--- HERE
                            training, momentum, eps, activation, activation_param, None)
'inplace_abn' is being compiled since it was called from 'InplaceAbn.forward'
  File "/workspace/timm-models/pytorch-image-models/timm/models/layers/inplace_abn.py", line 82
    def forward(self, x):
        output = inplace_abn(
        ~~~~~~~~~~~~~~~~~~~~~
            x, self.weight, self.bias, self.running_mean, self.running_var,
            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
            self.training, self.momentum, self.eps, self.act_name, self.act_param)
            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        if isinstance(output, tuple):
            output = output[0]
" while running benchmark.
Traceback (most recent call last):
  File "benchmark.py", line 721, in <module>
    main()
  File "benchmark.py", line 705, in main
    results = benchmark(args)
  File "benchmark.py", line 637, in benchmark
    run_results = _try_run(
  File "benchmark.py", line 579, in _try_run
    bench = bench_fn(model_name=model_name, batch_size=batch_size, **bench_kwargs)
  File "benchmark.py", line 391, in __init__
    super().__init__(model_name=model_name, device=device, torchscript=torchscript, **kwargs)
  File "benchmark.py", line 266, in __init__
    self.model = torch.jit.script(self.model)
  File "/opt/pytorch/pytorch/torch/jit/_script.py", line 1286, in script
    return torch.jit._recursive.create_script_module(
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 476, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/opt/pytorch/pytorch/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/opt/pytorch/pytorch/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 538, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/opt/pytorch/pytorch/torch/jit/_script.py", line 615, in _construct
    init_fn(script_module)
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 516, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 542, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 393, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
  File "/opt/pytorch/pytorch/torch/jit/_recursive.py", line 863, in try_compile_fn
    return torch.jit.script(fn, _rcb=rcb)
  File "/opt/pytorch/pytorch/torch/jit/_script.py", line 1343, in script
    fn = torch._C._jit_script_compile(
RuntimeError: 
Python builtin <built-in method apply of FunctionMeta object at 0x55c069d6eaa0> is currently not supported in Torchscript:
  File "/opt/conda/lib/python3.8/site-packages/inplace_abn/functions.py", line 153
def inplace_abn(x, weight, bias, running_mean, running_var,
                training=True, momentum=0.1, eps=1e-05, activation="leaky_relu", activation_param=0.01):
    return InPlaceABN.apply(x, weight, bias, running_mean, running_var,
           ~~~~~~~~~~~~~~~~ <--- HERE
                            training, momentum, eps, activation, activation_param, None)
'inplace_abn' is being compiled since it was called from 'InplaceAbn.forward'
  File "/workspace/timm-models/pytorch-image-models/timm/models/layers/inplace_abn.py", line 82
    def forward(self, x):
        output = inplace_abn(
        ~~~~~~~~~~~~~~~~~~~~~
            x, self.weight, self.bias, self.running_mean, self.running_var,
            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
            self.training, self.momentum, self.eps, self.act_name, self.act_param)
            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        if isinstance(output, tuple):
            output = output[0]

Versions

nightly core build 5589033