logo
Browse Source

Fix onnx export device issue

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 1 year ago
parent
commit
158d86d00e
  1. 2
      timm_image.py

2
timm_image.py

@ -166,7 +166,7 @@ class TimmImage(NNOperator):
path = path + '.onnx'
else:
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':
torch.save(self._model, path)
elif format == 'torchscript':

Loading…
Cancel
Save