Browse Source
Fix onnx export
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
1 additions and
1 deletions
-
auto_transformers.py
|
|
@ -188,7 +188,7 @@ class AutoTransformers(NNOperator): |
|
|
raise AttributeError('Unsupported model_type.') |
|
|
raise AttributeError('Unsupported model_type.') |
|
|
|
|
|
|
|
|
dummy_input = 'test sentence' |
|
|
dummy_input = 'test sentence' |
|
|
inputs = self.tokenizer(dummy_input, padding=True, truncation=True, return_tensors='pt') # a dictionary |
|
|
|
|
|
|
|
|
inputs = self.tokenizer(dummy_input, padding=True, truncation=True, return_tensors='pt').to(self.device) |
|
|
if model_type == 'pytorch': |
|
|
if model_type == 'pytorch': |
|
|
torch.save(self._model, output_file) |
|
|
torch.save(self._model, output_file) |
|
|
elif model_type == 'torchscript': |
|
|
elif model_type == 'torchscript': |
|
|
|
