specify shrinkage as run options not session options

qdrant · Mar 4, 2025 · 5c46b17 · 5c46b17
1 parent f63333d
commit 5c46b17
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py
@@ -84,7 +84,7 @@ def onnx_embed(
         onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
 
         run_options = ort.RunOptions()
-        run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "0")
+        run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", "gpu:0")
         model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
         return OnnxOutputContext(
             model_output=model_output[0],