Browse Source
Fix onnx export device issue
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
1 additions and
1 deletions
-
timm_image.py
|
|
@ -166,7 +166,7 @@ class TimmImage(NNOperator): |
|
|
path = path + '.onnx' |
|
|
path = path + '.onnx' |
|
|
else: |
|
|
else: |
|
|
raise AttributeError(f'Invalid format {format}.') |
|
|
raise AttributeError(f'Invalid format {format}.') |
|
|
dummy_input = torch.rand((1,) + self.config['input_size']) |
|
|
|
|
|
|
|
|
dummy_input = torch.rand((1,) + self.config['input_size']).to(self.device) |
|
|
if format == 'pytorch': |
|
|
if format == 'pytorch': |
|
|
torch.save(self._model, path) |
|
|
torch.save(self._model, path) |
|
|
elif format == 'torchscript': |
|
|
elif format == 'torchscript': |
|
|
|
