From 158d86d00e79020b8f805d0cb62c30b10cf51602 Mon Sep 17 00:00:00 2001
From: Jael Gu <mengjia.gu@zilliz.com>
Date: Fri, 3 Feb 2023 13:07:29 +0800
Subject: [PATCH] Fix onnx export device issue

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

diff --git a/timm_image.py b/timm_image.py
index 45dca20..f59c7bd 100644
--- a/timm_image.py
+++ b/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':