diff options
Diffstat (limited to 'model-integration/src')
-rw-r--r-- | model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluator.java | 13 |
1 files changed, 12 insertions, 1 deletions
diff --git a/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluator.java b/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluator.java index 9961c24005c..66eb8caabd0 100644 --- a/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluator.java +++ b/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluator.java @@ -20,7 +20,7 @@ import java.util.Map; * * @author lesters */ -public class OnnxEvaluator { +public class OnnxEvaluator implements AutoCloseable { private final OrtEnvironment environment; private final OrtSession session; @@ -88,6 +88,17 @@ public class OnnxEvaluator { } } + @Override + public void close() throws IllegalStateException { + try { + session.close(); + } catch (OrtException e) { + throw new IllegalStateException("Failed to close ONNX session", e); + } catch (IllegalStateException e) { + throw new IllegalStateException("Already closed", e); + } + } + private static OrtSession createSession(String modelPath, OrtEnvironment environment, OnnxEvaluatorOptions options, boolean tryCuda) { if (options == null) { options = new OnnxEvaluatorOptions(); |