add close method to all nd manager (#3225)

* add close method to all nd manager
deepjavalibrary · May 31, 2024 · 8955f6f · 8955f6f
1 parent e939474
commit 8955f6f
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDManager.java b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDManager.java
@@ -104,6 +104,16 @@ public NDList load(Path path) {
         return new NDList(new LgbmDataset(this, null, path));
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void close() {
+        super.close();
+        if (alternativeManager != null) {
+            alternativeManager.close();
+            alternativeManager = null;
+        }
+    }
+
     /** The SystemManager is the root {@link LgbmNDManager} of which all others are children. */
     private static final class SystemManager extends LgbmNDManager implements SystemNDManager {
 

diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java
@@ -180,6 +180,16 @@ public NDArray createCSR(Buffer buffer, long[] indptr, long[] indices, Shape sha
         return new XgbNDArray(this, alternativeManager, handle, shape, SparseFormat.CSR);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void close() {
+        super.close();
+        if (alternativeManager != null) {
+            alternativeManager.close();
+            alternativeManager = null;
+        }
+    }
+
     /** The SystemManager is the root {@link XgbNDManager} of which all others are children. */
     private static final class SystemManager extends XgbNDManager implements SystemNDManager {
 

diff --git a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDManager.java b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDManager.java
@@ -126,6 +126,16 @@ public NDArray ones(Shape shape, DataType dataType) {
         return create(bb, shape, dataType);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void close() {
+        super.close();
+        if (alternativeManager != null) {
+            alternativeManager.close();
+            alternativeManager = null;
+        }
+    }
+
     /** The SystemManager is the root {@link TrtNDManager} of which all others are children. */
     private static final class SystemManager extends TrtNDManager implements SystemNDManager {