Fix onnx export

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
2 years ago · e1a1abd518
1 changed files with 1 additions and 1 deletions
--- a/auto_transformers.py
+++ b/auto_transformers.py
@ -188,7 +188,7 @@ class AutoTransformers(NNOperator):
                raise AttributeError('Unsupported model_type.')

        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':
            torch.save(self._model, output_file)
        elif model_type == 'torchscript':