From b2337e8cd403649248831e35734e189dc17bdb3b Mon Sep 17 00:00:00 2001
From: Jael Gu <mengjia.gu@zilliz.com>
Date: Tue, 7 Feb 2023 18:45:39 +0800
Subject: [PATCH] Optimize triton

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
---
 auto_transformers.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/auto_transformers.py b/auto_transformers.py
index 911319c..b4055ce 100644
--- a/auto_transformers.py
+++ b/auto_transformers.py
@@ -128,12 +128,12 @@ class AutoTransformers(NNOperator):
             log.error(f'Fail to tokenize inputs: {e}')
             raise e
         try:
-            outs = self.model(**inputs)
+            outs = self.model(**inputs).to('cpu')
         except Exception as e:
             log.error(f'Invalid input for the model: {self.model_name}')
             raise e
         outs = self.post_proc(outs, inputs)
-        features = outs.cpu().detach().numpy()
+        features = outs.detach().numpy()
         if isinstance(data, str):
             features = features.squeeze(0)
         else: