Fix device issue for onnx export

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
2 years ago · 91b00e64fa
1 changed files with 1 additions and 1 deletions
--- a/isc.py
+++ b/isc.py
@ -130,7 +130,7 @@ class Isc(NNOperator):
                path = path + '.onnx'
            else:
                raise ValueError(f'Invalid format {format}.')
-        dummy_input = torch.rand(1, 3, 224, 224)
+        dummy_input = torch.rand(1, 3, 224, 224).to(self.device)
        if format == 'pytorch':
            torch.save(self._model, path)
        elif format == 'torchscript':