Spaces:
Running
Running
Update app.py
Browse files
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 |
-
|
109 |
-
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
|
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 |
|