How to convert layer_norm layer to ONNX?

I’m trying to convert my model to ONNX format for further deployment in TensorRT. Here is a sample code to illustrate my problem in layer_norm here.

import torch
from torch import nn

class ExportModel(nn.Module):
    def __init__(self):

    def forward(self, x):
        # n, c, h, w = x.shape
        # y = nn.functional.layer_norm(x, [c, h, w]])       # not working
        # y = nn.functional.layer_norm(x, x.size()[1:])     # not working
        y = nn.functional.layer_norm(x, [16, 32, 128])

        return y

def main():
    model = ExportModel()

    dummy_input = torch.randn(64, 16, 32, 128)
    input_names = [ "input" ]
    output_names = [ "output" ]

    with torch.no_grad():
            model, dummy_input, "sample.onnx", verbose=True,
            input_names=input_names, output_names=output_names

if __name__ == '__main__':

It could only work when the parameter of layer_norm is constant number. If not, the following error will occur.

Traceback (most recent call last):
  File "", line 31, in <module>
  File "", line 26, in main
    verbose=True, input_names=input_names, output_names=output_names
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 148, in export
    strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 66, in export
    dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 409, in _export
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 289, in _model_to_graph
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 132, in _optimize_graph
    graph = torch._C._jit_pass_onnx(graph, operator_export_type)
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 179, in _run_symbolic_function
    return utils._run_symbolic_function(*args, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 647, in _run_symbolic_function
    return op_fn(g, *inputs, **attrs)
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 128, in wrapper
    args = [_parse_arg(arg, arg_desc) for arg, arg_desc in zip(args, arg_descriptors)]
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 128, in <listcomp>
    args = [_parse_arg(arg, arg_desc) for arg, arg_desc in zip(args, arg_descriptors)]
  File "/opt/conda/lib/python3.6/site-packages/torch/onnx/", line 81, in _parse_arg
    "', since it's not constant, please try to make "
RuntimeError: Failed to export an ONNX attribute 'onnx::Gather', since it's not constant, please try to make things (e.g., kernel size) static if possible

I have few code blocks in my model have layer_norm op. It would turn into some ugly code if I explicitly mark all parameters constant number. Is there any “best practice” of how to use dynamic shape for this kind of use case?

Thanks in advance.

Try this codes. it converts tensor variables to integer ones.
I confirmed that it works for your example.
n, c, h, w = x.shape
c1 = c.item()
h1 = h.item()
w1 = w.item()
y = nn.functional.layer_norm(x, [c1, h1, w1])