From 884323d93d7cb1e42d7bfb6a62395916ab081ff0 Mon Sep 17 00:00:00 2001
From: ChengZi <chen.zhang@zilliz.com>
Date: Thu, 14 Sep 2023 19:16:54 +0800
Subject: [PATCH] export onnx using cpu model

Signed-off-by: ChengZi <chen.zhang@zilliz.com>
---
 rerank.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rerank.py b/rerank.py
index 4840319..cb663c2 100644
--- a/rerank.py
+++ b/rerank.py
@@ -122,7 +122,7 @@ class ReRank(NNOperator):
             onnx_config = model_onnx_config(self._model.config)
             onnx_inputs, onnx_outputs = export(
                 self.tokenizer,
-                self._model,
+                self._model.to('cpu'),
                 config=onnx_config,
                 opset=13,
                 output=Path(path)