wasmdashai commited on
Commit
3442992
·
verified ·
1 Parent(s): 3e22085

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -9
app.py CHANGED
@@ -104,15 +104,13 @@ class OnnxModelConverter:
104
  # Trace the decoder part of the model
105
  example_inputs = inputs.input_ids.type(torch.LongTensor)
106
 
107
- torch.onnx.export(
108
- model,
109
- example_inputs,
110
- onnx_file,
111
- opset_version=11,
112
- input_names=['input'],
113
- output_names=['output'],
114
- dynamic_axes={'input': {0: 'batch_size', 1: 'sequence_length'}, 'output': {0: 'batch_size'}}
115
- )
116
 
117
  return self.download_file(onnx_file)
118
 
 
104
  # Trace the decoder part of the model
105
  example_inputs = inputs.input_ids.type(torch.LongTensor)
106
 
107
+ torch.onnx.export(model,
108
+ example_inputs,
109
+ onnx_file,
110
+ input_names=["input"],
111
+ output_names=["output"],
112
+ dynamic_axes={"input": {0: "batch_size", 1: "sequence_length"},
113
+ "output": {0: "batch_size", 1: "sequence_length"}})
 
 
114
 
115
  return self.download_file(onnx_file)
116