[ort] Add support for registerCustomOpLibrary. (#1771)

Change-Id: Ic3ec4249adb432b2a54668b7a45481698fa83ea5
deepjavalibrary · Jul 6, 2022 · 7cecdfd · 7cecdfd
1 parent a7ee401
commit 7cecdfd
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
@@ -179,6 +179,11 @@ private SessionOptions getSessionOptions(Map<String, ?> options) throws OrtExcep
             ortSession.setCPUArenaAllocator(true);
         }
 
+        String customOpLibrary = (String) options.get("customOpLibrary");
+        if (customOpLibrary != null) {
+            ortSession.registerCustomOpLibrary(customOpLibrary);
+        }
+
         Device device = manager.getDevice();
         if (options.containsKey("ortDevice")) {
             String ortDevice = (String) options.get("ortDevice");