[ONNX Runtime] Fix OrtNDArray double close issue (deepjavalibrary#1861)

patins1 · Aug 26, 2022 · abe9b0f · abe9b0f
1 parent 5a43375
commit abe9b0f
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/...es/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java b/...es/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
@@ -76,7 +76,12 @@ public Shape getShape() {
     /** {@inheritDoc} */
     @Override
     public void intern(NDArray replaced) {
-        tensor.getAndSet(((OrtNDArray) replaced).getTensor()).close();
+        OrtNDArray arr = (OrtNDArray) replaced;
+        OnnxTensor oldHandle = tensor.getAndSet(arr.tensor.getAndSet(null));
+        if (oldHandle != null) {
+            oldHandle.close();
+        }
+        replaced.close();
     }
 
     /** {@inheritDoc} */
@@ -113,10 +118,10 @@ public ByteBuffer toByteBuffer() {
     /** {@inheritDoc} */
     @Override
     public void close() {
-        super.close();
         OnnxTensor ortTensor = tensor.getAndSet(null);
         if (ortTensor != null) {
             ortTensor.close();
         }
+        super.close();
     }
 }