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 index b9f79a124ee..80351ee79f1 100644 --- 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 index 81f9708e72b..d1406b14011 100644 --- 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 index f69276b61a8..ebb4de96ca3 100644 --- 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 {