Fix capped manager on Tf engine (#1952)

* fix_capped_manager_on_Tf_engine
deepjavalibrary · Aug 25, 2022 · e21b1ac · e21b1ac
1 parent 085741a
commit e21b1ac
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDArray.java b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDArray.java
@@ -55,6 +55,14 @@ public class LgbmNDArray extends NDArrayAdapter {
  handle = new AtomicReference<>();
  }
 
+ /** {@inheritDoc} */
+ @Override
+ public void returnResource(NDManager manager) {
+ detach();
+ this.manager = manager;
+ manager.attachUncappedInternal(getUid(), this);
+ }
+
  /**
  * Returns the native LightGBM handle to the array.
  *

diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java
@@ -98,6 +98,14 @@ public void intern(NDArray replaced) {
  format = array.format;
  }
 
+ /** {@inheritDoc} */
+ @Override
+ public void returnResource(NDManager manager) {
+ detach();
+ this.manager = manager;
+ manager.attachUncappedInternal(getUid(), this);
+ }
+
  /** {@inheritDoc} */
  @Override
  public void detach() {

diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
@@ -228,6 +228,14 @@ public void attach(NDManager manager) {
  manager.attachInternal(getUid(), this);
  }
 
+ /** {@inheritDoc} */
+ @Override
+ public void returnResource(NDManager manager) {
+ detach();
+ this.manager = (TfNDManager) manager;
+ manager.attachUncappedInternal(getUid(), this);
+ }
+
  /** {@inheritDoc} */
  @Override
  public void tempAttach(NDManager manager) {