ONNX export hangs

[genolve]

I’m trying to export the model to ONNX and the export simply hangs. Before diving into debugging, wanted check if anyone has had success with an ONNX export.

I have torch version: 1.13.0.dev20220616

Export command:

import torch.utils.model_zoo as model_zoo
import torch.onnx
import netron
device =  torch.device('cpu')
torch.onnx.enable_log() 
torch.set_num_threads(1)
#The tuple should contain model inputs such that model(*args) is a valid invocation of the model
tinput=(torch.tensor(firstbatch['input_ids']).to(device)
       )
model.cpu()
model.eval()
torch.onnx.export(model,                     # model being run  debug: export_to_pretty_string
                  tinput,                    # model input (or a tuple for multiple inputs)
                  dataset+"_routing.onnx",   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=16,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names  = ['input_ids'],   # the model's input names
                  output_names = ['output'],  # the model's output names
                  verbose=True
                  #,operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK
                 )

[genolve]

I ended up switching to x-transformer which exports to ONNX with minimal modifications. It also hangs on the first export attempt; just interrupt the kernel and the second export works.