ONNX Conversion for SQuAD example

onnx_convert.py

torch.onnx.export(model = torch_model,
                 args = (dummy_input['input_ids'].to(device), dummy_input['attention_mask'].to(device)), #Tuple of inputs for models that take more than one input
                 f = onnx_path,
                 export_params = True,
                 opset_version=11,
                 do_constant_folding = True,
                 input_names = ['input_ids', 'attention_mask'],
                 output_names = ['output'],
                 dynamic_axes = {'input_ids': {0: 'batch_size'},
                                'attention_mask': {0: 'batch_size'},
                                'output': {0: 'batch_size'}}
                 )