refactor: remove inference model pytorch from onnx (#661)

jina-ai · May 4, 2022 · 2c1f32a · 2c1f32a
1 parent 665deb3
commit 2c1f32a
Showing 1 changed file with 9 additions and 10 deletions.
diff --git a/server/clip_server/executors/clip_onnx.py b/server/clip_server/executors/clip_onnx.py
@@ -55,17 +55,16 @@ async def encode(self, docs: 'DocumentArray', **kwargs):
         _img_da = docs.find({'blob': {'$exists': True}})
         _txt_da = docs.find({'text': {'$exists': True}})
 
-        with torch.inference_mode():
-            # for image
-            if _img_da:
-                _img_da.apply(self._preproc_image)
-                _img_da.embeddings = self._model.encode_image(_img_da.tensors)
+        # for image
+        if _img_da:
+            _img_da.apply(self._preproc_image)
+            _img_da.embeddings = self._model.encode_image(_img_da.tensors)
 
-            # for text
-            if _txt_da:
-                texts = self._preproc_text(_txt_da)
-                _txt_da.embeddings = self._model.encode_text(_txt_da.tensors)
-                _txt_da.texts = texts
+        # for text
+        if _txt_da:
+            texts = self._preproc_text(_txt_da)
+            _txt_da.embeddings = self._model.encode_text(_txt_da.tensors)
+            _txt_da.texts = texts
 
         # drop tensors
         docs.tensors = None