Feature Extraction
sentence-transformers
ONNX
English
sentence-similarity
Inference Endpoints
shuttie commited on
Commit
3daf9af
1 Parent(s): 7176493

set proper output layer name

Browse files
Files changed (2) hide show
  1. convert.py +1 -1
  2. pytorch_model.onnx +2 -2
convert.py CHANGED
@@ -13,6 +13,6 @@ inputs = {"input_ids": torch.ones(1, max_seq_length, dtype=torch.int64),
13
  symbolic_names = {0: 'batch_size', 1: 'max_seq_len'}
14
 
15
  torch.onnx.export(model, args=tuple(inputs.values()), f='pytorch_model.onnx', export_params=True,
16
- input_names=['input_ids', 'attention_mask', 'token_type_ids'], output_names=['output'],
17
  dynamic_axes={'input_ids': symbolic_names, 'attention_mask': symbolic_names, 'token_type_ids': symbolic_names})
18
 
 
13
  symbolic_names = {0: 'batch_size', 1: 'max_seq_len'}
14
 
15
  torch.onnx.export(model, args=tuple(inputs.values()), f='pytorch_model.onnx', export_params=True,
16
+ input_names=['input_ids', 'attention_mask', 'token_type_ids'], output_names=['last_hidden_state'],
17
  dynamic_axes={'input_ids': symbolic_names, 'attention_mask': symbolic_names, 'token_type_ids': symbolic_names})
18
 
pytorch_model.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:79c87071eeceb3a40ef53b96aa8b3f9bfd4d3152023a5ae25ed134eed59c724e
3
- size 90984197
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:47206d2e5b865151da68245745c426ebc496d21200bbd73154c3a2ddc699b0a1
3
+ size 90984263