ONNX.py

model.eval()
dummy_input = torch.randn(1, 3, INPUT_HEIGHT, INPUT_WIDTH, 
                          dtype=torch.float32, device=DEVICE)

output = model(dummy_input.detach())

input_names = ['data']
output_names = ['output1']

torch.onnx.export(model, dummy_input, 
                  ONNX_PATH, verbose=False, 
                  input_names=input_names, output_names=output_names)