From 91b00e64faf02a17a51ebc35e1bd5eb7e7329457 Mon Sep 17 00:00:00 2001
From: Jael Gu <mengjia.gu@zilliz.com>
Date: Fri, 3 Feb 2023 14:59:47 +0800
Subject: [PATCH] Fix device issue for onnx export

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
---
 isc.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/isc.py b/isc.py
index 8a5ee07..52dcfc5 100644
--- 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':