diff options
author | Bjørn Christian Seime <bjorncs@yahooinc.com> | 2023-02-22 11:23:49 +0100 |
---|---|---|
committer | Bjørn Christian Seime <bjorncs@yahooinc.com> | 2023-02-22 11:23:49 +0100 |
commit | c5513d25475c78ce6a3ecd5e03b278f3eebca481 (patch) | |
tree | a5aecbcc19ac7415401b83aa0619f1f89b2622a0 /model-integration | |
parent | c744731f33372b8176f09f21959f4a2b321f4b64 (diff) |
Make OnnxEvaluator closable
Diffstat (limited to 'model-integration')
-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(); |