Can not export Huggingface GPT-2 to ONNX

Hi, thank you for your great framework.

I am trying to convert Huggingface GPT-2 model to ONNX format using this script but it doesn't work.

Here is the error:

/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py:100: UserWarning: `example_outputs' is deprecated and ignored. Will be removed in next PyTorch release.
  warnings.warn("`example_outputs' is deprecated and ignored. Will be removed in "
/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py:1297: UserWarning: No names were found for specified dynamic axes of provided input.Automatically generated names will be applied to each dynamic axes of input input1
  "Automatically generated names will be applied to each dynamic axes of input {}".format(key))
/usr/local/lib/python3.7/dist-packages/transformers/models/gpt2/modeling_gpt2.py:196: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  attn_weights = attn_weights / (float(value.size(-1)) ** 0.5)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-1-150445c8bc3a> in <module>()
    173 
    174 
--> 175 gpt2_test()

6 frames
<ipython-input-1-150445c8bc3a> in gpt2_test()
    137                                          opset_version=10,
    138                                          input_names=['input1'],
--> 139                                          dynamic_axes={'input1': [0, 1, 2, 3]})
    140 
    141         # Test exported model with TensorProto data saved in files

<ipython-input-1-150445c8bc3a> in save_model(name, model, inputs, outputs, input_names, output_names, **kwargs)
     88     else:
     89         torch.onnx.export(model, inputs, model_dir, verbose=True, input_names=input_names,
---> 90                           output_names=output_names, example_outputs=outputs, **kwargs)
     91 
     92     test_data_dir = os.path.join(dir, data_dir)

/usr/local/lib/python3.7/dist-packages/torch/onnx/__init__.py in export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, _retain_param_name, do_constant_folding, example_outputs, strip_doc_string, dynamic_axes, keep_initializers_as_inputs, custom_opsets, enable_onnx_checker, use_external_data_format)
    318                         _retain_param_name, do_constant_folding, example_outputs,
    319                         strip_doc_string, dynamic_axes, keep_initializers_as_inputs,
--> 320                         custom_opsets, enable_onnx_checker, use_external_data_format)
    321 
    322 

/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py in export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, opset_version, _retain_param_name, do_constant_folding, example_outputs, strip_doc_string, dynamic_axes, keep_initializers_as_inputs, custom_opsets, enable_onnx_checker, use_external_data_format)
    109             do_constant_folding=do_constant_folding, example_outputs=example_outputs,
    110             dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs,
--> 111             custom_opsets=custom_opsets, use_external_data_format=use_external_data_format)
    112 
    113 

/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, example_outputs, opset_version, do_constant_folding, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size, custom_opsets, add_node_names, use_external_data_format, onnx_shape_inference)
    727                                 fixed_batch_size=fixed_batch_size,
    728                                 training=training,
--> 729                                 dynamic_axes=dynamic_axes)
    730 
    731             # TODO: Don't allocate a in-memory string for the protobuf

/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py in _model_to_graph(model, args, verbose, input_names, output_names, operator_export_type, example_outputs, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes)
    499                             fixed_batch_size=fixed_batch_size, params_dict=params_dict,
    500                             dynamic_axes=dynamic_axes, input_names=input_names,
--> 501                             module=module)
    502     from torch.onnx.symbolic_helper import _onnx_shape_inference
    503     if isinstance(model, torch.jit.ScriptModule) or isinstance(model, torch.jit.ScriptFunction):

/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py in _optimize_graph(graph, operator_export_type, _disable_torch_constant_prop, fixed_batch_size, params_dict, dynamic_axes, input_names, module)
    213         input_names = [] if input_names is None else input_names
    214         dynamic_axes = {} if dynamic_axes is None else dynamic_axes
--> 215         torch._C._jit_pass_onnx_set_dynamic_input_shape(graph, dynamic_axes, input_names)
    216     graph = torch._C._jit_pass_onnx(graph, operator_export_type)
    217     torch._C._jit_pass_lint(graph)

RuntimeError: Dynamic shape axis should be no more than the shape dimension for input1_dynamic_axes_4

Here is the Colab notebook that I used to run the code: https://colab.research.google.com/drive/1j8366ADwIZv-GPjzu9NLSvs6gzQP_CpT?usp=sharing

Thank you so much!

Edit: I tried both onnx current newest version and version 1.6, but they give the same error.

Open the linked GPT2-export.py script and reproduce it in the linked Colab notebook. Start at the torch.onnx.export call and its dynamic_axes argument, then compare it with the reported input shape. Done means the GPT-2 model exports without the dynamic-axis error and the script can test the resulting ONNX model.