Repository: deepjavalibrary/djl
Pull request #1252 (2021-09-26 03:15:31)

Fix model close() NPE issue
Change-Id: I5ff882caa1c24a71fbc282bb638a3e5736a77408

## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1252?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1252](https://codecov.io/gh/deepjavalibrary/djl/pull/1252?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (b33ee44) into [master](https://codecov.io/gh/deepjavalibrary/djl/commit/5907d505a55351394ca26babcd0d4dd7411a500a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (5907d50) will **decrease** coverage by `0.03%`.
> The diff coverage is `76.92%`.

[![Impacted file tree graph](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/graphs/tree.svg?width=650&height=150&src=pr&token=mDKVNwFXRV&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)](https://codecov.io/gh/deepjavalibrary/djl/pull/1252?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

```diff
@@             Coverage Diff              @@
##             master    #1252      +/-   ##
============================================
- Coverage     70.23%   70.19%   -0.04%     
- Complexity     5063     5066       +3     
============================================
  Files           485      485              
  Lines         22262    22271       +9     
  Branches       2375     2375              
============================================
- Hits          15635    15634       -1     
- Misses         5414     5426      +12     
+ Partials       1213     1211       -2     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1252?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [...src/main/java/ai/djl/tensorrt/engine/TrtModel.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy90ZW5zb3JydC9zcmMvbWFpbi9qYXZhL2FpL2RqbC90ZW5zb3JydC9lbmdpbmUvVHJ0TW9kZWwuamF2YQ==) | `0.00% <0.00%> (ø)` | |
| [.../main/java/ai/djl/paddlepaddle/engine/PpModel.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy9wYWRkbGVwYWRkbGUvcGFkZGxlcGFkZGxlLWVuZ2luZS9zcmMvbWFpbi9qYXZhL2FpL2RqbC9wYWRkbGVwYWRkbGUvZW5naW5lL1BwTW9kZWwuamF2YQ==) | `60.41% <66.66%> (-0.46%)` | :arrow_down: |
| [...rc/main/java/ai/djl/tflite/engine/TfLiteModel.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy90ZmxpdGUvdGZsaXRlLWVuZ2luZS9zcmMvbWFpbi9qYXZhL2FpL2RqbC90ZmxpdGUvZW5naW5lL1RmTGl0ZU1vZGVsLmphdmE=) | `47.50% <66.66%> (+0.13%)` | :arrow_up: |
| [...oost/src/main/java/ai/djl/ml/xgboost/XgbModel.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy9tbC94Z2Jvb3N0L3NyYy9tYWluL2phdmEvYWkvZGpsL21sL3hnYm9vc3QvWGdiTW9kZWwuamF2YQ==) | `36.17% <100.00%> (+2.83%)` | :arrow_up: |
| [...rc/main/java/ai/djl/tensorflow/engine/TfModel.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy90ZW5zb3JmbG93L3RlbnNvcmZsb3ctZW5naW5lL3NyYy9tYWluL2phdmEvYWkvZGpsL3RlbnNvcmZsb3cvZW5naW5lL1RmTW9kZWwuamF2YQ==) | `44.70% <100.00%> (+1.33%)` | :arrow_up: |
| [...gine/src/main/java/ai/djl/dlr/engine/DlrModel.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy9kbHIvZGxyLWVuZ2luZS9zcmMvbWFpbi9qYXZhL2FpL2RqbC9kbHIvZW5naW5lL0Rsck1vZGVsLmphdmE=) | `51.61% <0.00%> (-3.23%)` | :arrow_down: |
| [api/src/main/java/ai/djl/inference/Predictor.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL2luZmVyZW5jZS9QcmVkaWN0b3IuamF2YQ==) | `88.46% <0.00%> (-2.89%)` | :arrow_down: |
| [...rt/src/main/java/ai/djl/tensorrt/jni/LibUtils.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy90ZW5zb3JydC9zcmMvbWFpbi9qYXZhL2FpL2RqbC90ZW5zb3JydC9qbmkvTGliVXRpbHMuamF2YQ==) | `71.05% <0.00%> (-2.64%)` | :arrow_down: |
| [api/src/main/java/ai/djl/util/Platform.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL3V0aWwvUGxhdGZvcm0uamF2YQ==) | `69.81% <0.00%> (-1.89%)` | :arrow_down: |
| [...ain/java/ai/djl/tensorflow/engine/TfNDManager.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy90ZW5zb3JmbG93L3RlbnNvcmZsb3ctZW5naW5lL3NyYy9tYWluL2phdmEvYWkvZGpsL3RlbnNvcmZsb3cvZW5naW5lL1RmTkRNYW5hZ2VyLmphdmE=) | `85.71% <0.00%> (-1.69%)` | :arrow_down: |
| ... and [5 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1252/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1252?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1252?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [5907d50...b33ee44](https://codecov.io/gh/deepjavalibrary/djl/pull/1252?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbModel.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbModel.java
index 365930ad7..3d323e290 100644
--- a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbModel.java
+++ b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbModel.java
@@ -101,7 +101,10 @@ public class XgbModel extends BaseModel {
     /** {@inheritDoc} */
     @Override
     public void close() {
-        ((XgbSymbolBlock) block).close();
+        if (block != null) {
+            ((XgbSymbolBlock) block).close();
+            block = null;
+        }
         super.close();
     }
 }
diff --git a/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpModel.java b/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpModel.java
index b0930a2f9..974a3e712 100644
--- a/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpModel.java
+++ b/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpModel.java
@@ -131,7 +131,10 @@ public class PpModel extends BaseModel {
     /** {@inheritDoc} */
     @Override
     public void close() {
-        JniUtils.deletePredictor(paddlePredictor);
+        if (paddlePredictor != null) {
+            JniUtils.deletePredictor(paddlePredictor);
+            paddlePredictor = null;
+        }
         super.close();
     }
 }
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfModel.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfModel.java
index 989e1c40c..e645eb463 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfModel.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfModel.java
@@ -185,7 +185,10 @@ public class TfModel extends BaseModel {
     /** {@inheritDoc} */
     @Override
     public void close() {
-        ((TfSymbolBlock) block).close();
+        if (block != null) {
+            ((TfSymbolBlock) block).close();
+            block = null;
+        }
         super.close();
     }
 }
diff --git a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtModel.java b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtModel.java
index 8f66932b2..b65b61950 100644
--- a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtModel.java
+++ b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtModel.java
@@ -128,6 +128,7 @@ public class TrtModel extends BaseModel {
     public void close() {
         if (block != null) {
             ((TrtSymbolBlock) block).close();
+            block = null;
         }
         super.close();
     }
diff --git a/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/engine/TfLiteModel.java b/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/engine/TfLiteModel.java
index 1b92bf0c9..fffcba7cb 100644
--- a/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/engine/TfLiteModel.java
+++ b/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/engine/TfLiteModel.java
@@ -99,7 +99,10 @@ public class TfLiteModel extends BaseModel {
     /** {@inheritDoc} */
     @Override
     public void close() {
-        ((TfLiteSymbolBlock) block).close();
+        if (block != null) {
+            ((TfLiteSymbolBlock) block).close();
+            block = null;
+        }
         super.close();
     }
 }


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/AlexNetTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/AlexNetTest.java
index cc9ff08df..cd1012937 100644
--- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/AlexNetTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/AlexNetTest.java
@@ -153,48 +153,47 @@ public class AlexNetTest {
 
     @Test
     public void testOutputShapes() {
-        NDManager manager = NDManager.newBaseManager();
-        int batchSize = 2;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
-        Shape currentShape = x.getShape();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            int batchSize = 2;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
+            Shape currentShape = x.getShape();
+
+            Block alexNet = AlexNet.builder().build();
+            alexNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            alexNet.initialize(manager, DataType.FLOAT32, currentShape);
+
+            Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
+            for (int i = 0; i < alexNet.getChildren().size(); i++) {
+
+                Shape[] newShape =
+                        alexNet.getChildren()
+                                .get(i)
+                                .getValue()
+                                .getOutputShapes(new Shape[] {currentShape});
+                currentShape = newShape[0];
+                shapeMap.put(alexNet.getChildren().get(i).getKey(), currentShape);
+            }
 
-        Block alexNet = AlexNet.builder().build();
-        alexNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        alexNet.initialize(manager, DataType.FLOAT32, currentShape);
-
-        Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
-        for (int i = 0; i < alexNet.getChildren().size(); i++) {
-
-            Shape[] newShape =
-                    alexNet.getChildren()
-                            .get(i)
-                            .getValue()
-                            .getOutputShapes(new Shape[] {currentShape});
-            currentShape = newShape[0];
-            shapeMap.put(alexNet.getChildren().get(i).getKey(), currentShape);
+            Assert.assertEquals(shapeMap.get("01Conv2d"), new Shape(batchSize, 96, 54, 54));
+            Assert.assertEquals(shapeMap.get("04Conv2d"), new Shape(batchSize, 256, 26, 26));
+            Assert.assertEquals(shapeMap.get("07Conv2d"), new Shape(batchSize, 384, 12, 12));
+            Assert.assertEquals(shapeMap.get("13LambdaBlock"), new Shape(batchSize, 256, 5, 5));
+            Assert.assertEquals(shapeMap.get("17Dropout"), new Shape(batchSize, 4096));
         }
-
-        Assert.assertEquals(shapeMap.get("01Conv2d"), new Shape(batchSize, 96, 54, 54));
-        Assert.assertEquals(shapeMap.get("04Conv2d"), new Shape(batchSize, 256, 26, 26));
-        Assert.assertEquals(shapeMap.get("07Conv2d"), new Shape(batchSize, 384, 12, 12));
-        Assert.assertEquals(shapeMap.get("13LambdaBlock"), new Shape(batchSize, 256, 5, 5));
-        Assert.assertEquals(shapeMap.get("17Dropout"), new Shape(batchSize, 4096));
-        manager.close();
     }
 
     @Test
     public void testForwardMethod() {
-        NDManager manager = NDManager.newBaseManager();
-        Block alexNet = AlexNet.builder().build();
-        int batchSize = 1;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
-        alexNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        alexNet.initialize(manager, DataType.FLOAT32, x.getShape());
-        NDArray xHat =
-                alexNet.forward(new ParameterStore(manager, true), new NDList(x), false)
-                        .singletonOrThrow();
-
-        Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
-        manager.close();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            Block alexNet = AlexNet.builder().build();
+            int batchSize = 1;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
+            alexNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            alexNet.initialize(manager, DataType.FLOAT32, x.getShape());
+            ParameterStore ps = new ParameterStore(manager, true);
+            NDArray xHat = alexNet.forward(ps, new NDList(x), false).singletonOrThrow();
+
+            Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
+        }
     }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/GoogLeNetTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/GoogLeNetTest.java
index 9d711ec60..748940d3f 100644
--- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/GoogLeNetTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/GoogLeNetTest.java
@@ -94,50 +94,49 @@ public class GoogLeNetTest {
 
     @Test
     public void testOutputShapes() {
-        NDManager manager = NDManager.newBaseManager();
-        int batchSize = 1;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 96, 96));
-        Shape currentShape = x.getShape();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            int batchSize = 1;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 96, 96));
+            Shape currentShape = x.getShape();
 
-        Block googLeNet = GoogLeNet.builder().build();
-        googLeNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        googLeNet.initialize(manager, DataType.FLOAT32, currentShape);
+            Block googLeNet = GoogLeNet.builder().build();
+            googLeNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            googLeNet.initialize(manager, DataType.FLOAT32, currentShape);
 
-        Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
-        for (int i = 0; i < googLeNet.getChildren().size(); i++) {
+            Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
+            for (int i = 0; i < googLeNet.getChildren().size(); i++) {
+                Shape[] newShape =
+                        googLeNet
+                                .getChildren()
+                                .get(i)
+                                .getValue()
+                                .getOutputShapes(new Shape[] {currentShape});
+                currentShape = newShape[0];
+                shapeMap.put(googLeNet.getChildren().get(i).getKey(), currentShape);
+            }
 
-            Shape[] newShape =
-                    googLeNet
-                            .getChildren()
-                            .get(i)
-                            .getValue()
-                            .getOutputShapes(new Shape[] {currentShape});
-            currentShape = newShape[0];
-            shapeMap.put(googLeNet.getChildren().get(i).getKey(), currentShape);
+            Assert.assertEquals(
+                    shapeMap.get("01SequentialBlock"), new Shape(batchSize, 64, 24, 24));
+            Assert.assertEquals(
+                    shapeMap.get("02SequentialBlock"), new Shape(batchSize, 192, 12, 12));
+            Assert.assertEquals(shapeMap.get("03SequentialBlock"), new Shape(batchSize, 480, 6, 6));
+            Assert.assertEquals(shapeMap.get("04SequentialBlock"), new Shape(batchSize, 832, 3, 3));
+            Assert.assertEquals(shapeMap.get("05SequentialBlock"), new Shape(batchSize, 1024));
         }
-
-        Assert.assertEquals(shapeMap.get("01SequentialBlock"), new Shape(batchSize, 64, 24, 24));
-        Assert.assertEquals(shapeMap.get("02SequentialBlock"), new Shape(batchSize, 192, 12, 12));
-        Assert.assertEquals(shapeMap.get("03SequentialBlock"), new Shape(batchSize, 480, 6, 6));
-        Assert.assertEquals(shapeMap.get("04SequentialBlock"), new Shape(batchSize, 832, 3, 3));
-        Assert.assertEquals(shapeMap.get("05SequentialBlock"), new Shape(batchSize, 1024));
-        manager.close();
     }
 
     @Test
     public void testForwardMethod() {
-        NDManager manager = NDManager.newBaseManager();
-        Block googLeNet = GoogLeNet.builder().build();
-        int batchSize = 1;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 28, 28));
-        googLeNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        googLeNet.initialize(manager, DataType.FLOAT32, x.getShape());
-        NDArray xHat =
-                googLeNet
-                        .forward(new ParameterStore(manager, true), new NDList(x), true)
-                        .singletonOrThrow();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            Block googLeNet = GoogLeNet.builder().build();
+            int batchSize = 1;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 28, 28));
+            googLeNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            googLeNet.initialize(manager, DataType.FLOAT32, x.getShape());
+            ParameterStore ps = new ParameterStore(manager, true);
+            NDArray xHat = googLeNet.forward(ps, new NDList(x), true).singletonOrThrow();
 
-        Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
-        manager.close();
+            Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
+        }
     }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/LeNetTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/LeNetTest.java
index 93346d3d9..edba40a52 100644
--- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/LeNetTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/LeNetTest.java
@@ -131,47 +131,45 @@ public class LeNetTest {
 
     @Test
     public void testOutputShapes() {
-        NDManager manager = NDManager.newBaseManager();
-        int batchSize = 1;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 28, 28));
-        Shape currentShape = x.getShape();
-
-        Block leNet = LeNet.builder().build();
-        leNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        leNet.initialize(manager, DataType.FLOAT32, currentShape);
-
-        Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
-        for (int i = 0; i < leNet.getChildren().size(); i++) {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            int batchSize = 1;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 28, 28));
+            Shape currentShape = x.getShape();
+
+            Block leNet = LeNet.builder().build();
+            leNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            leNet.initialize(manager, DataType.FLOAT32, currentShape);
+
+            Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
+            for (int i = 0; i < leNet.getChildren().size(); i++) {
+                Shape[] newShape =
+                        leNet.getChildren()
+                                .get(i)
+                                .getValue()
+                                .getOutputShapes(new Shape[] {currentShape});
+                currentShape = newShape[0];
+                shapeMap.put(leNet.getChildren().get(i).getKey(), currentShape);
+            }
 
-            Shape[] newShape =
-                    leNet.getChildren()
-                            .get(i)
-                            .getValue()
-                            .getOutputShapes(new Shape[] {currentShape});
-            currentShape = newShape[0];
-            shapeMap.put(leNet.getChildren().get(i).getKey(), currentShape);
+            Assert.assertEquals(shapeMap.get("01Conv2d"), new Shape(batchSize, 6, 28, 28));
+            Assert.assertEquals(shapeMap.get("04Conv2d"), new Shape(batchSize, 16, 10, 10));
+            Assert.assertEquals(shapeMap.get("08Linear"), new Shape(batchSize, 120));
+            Assert.assertEquals(shapeMap.get("12Linear"), new Shape(batchSize, 10));
         }
-
-        Assert.assertEquals(shapeMap.get("01Conv2d"), new Shape(batchSize, 6, 28, 28));
-        Assert.assertEquals(shapeMap.get("04Conv2d"), new Shape(batchSize, 16, 10, 10));
-        Assert.assertEquals(shapeMap.get("08Linear"), new Shape(batchSize, 120));
-        Assert.assertEquals(shapeMap.get("12Linear"), new Shape(batchSize, 10));
-        manager.close();
     }
 
     @Test
     public void testForwardMethod() {
-        NDManager manager = NDManager.newBaseManager();
-        Block leNet = LeNet.builder().build();
-        int batchSize = 1;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 28, 28));
-        leNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        leNet.initialize(manager, DataType.FLOAT32, x.getShape());
-        NDArray xHat =
-                leNet.forward(new ParameterStore(manager, true), new NDList(x), true)
-                        .singletonOrThrow();
-
-        Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
-        manager.close();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            Block leNet = LeNet.builder().build();
+            int batchSize = 1;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 28, 28));
+            leNet.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            leNet.initialize(manager, DataType.FLOAT32, x.getShape());
+            ParameterStore ps = new ParameterStore(manager, true);
+            NDArray xHat = leNet.forward(ps, new NDList(x), true).singletonOrThrow();
+
+            Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
+        }
     }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/NiNTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/NiNTest.java
index d2208429d..0d77cbb9f 100644
--- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/NiNTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/NiNTest.java
@@ -146,44 +146,48 @@ public class NiNTest {
 
     @Test
     public void testOutputShapes() {
-        NDManager manager = NDManager.newBaseManager();
-        int batchSize = 1;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
-        Shape currentShape = x.getShape();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            int batchSize = 1;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
+            Shape currentShape = x.getShape();
 
-        Block nin = NiN.builder().build();
-        nin.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        nin.initialize(manager, DataType.FLOAT32, currentShape);
+            Block nin = NiN.builder().build();
+            nin.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            nin.initialize(manager, DataType.FLOAT32, currentShape);
 
-        Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
-        for (int i = 0; i < nin.getChildren().size(); i++) {
+            Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
+            for (int i = 0; i < nin.getChildren().size(); i++) {
+                Shape[] newShape =
+                        nin.getChildren()
+                                .get(i)
+                                .getValue()
+                                .getOutputShapes(new Shape[] {currentShape});
+                currentShape = newShape[0];
+                shapeMap.put(nin.getChildren().get(i).getKey(), currentShape);
+            }
 
-            Shape[] newShape =
-                    nin.getChildren().get(i).getValue().getOutputShapes(new Shape[] {currentShape});
-            currentShape = newShape[0];
-            shapeMap.put(nin.getChildren().get(i).getKey(), currentShape);
+            Assert.assertEquals(
+                    shapeMap.get("01SequentialBlock"), new Shape(batchSize, 96, 54, 54));
+            Assert.assertEquals(
+                    shapeMap.get("03SequentialBlock"), new Shape(batchSize, 256, 26, 26));
+            Assert.assertEquals(
+                    shapeMap.get("05SequentialBlock"), new Shape(batchSize, 384, 12, 12));
+            Assert.assertEquals(shapeMap.get("08SequentialBlock"), new Shape(batchSize, 10, 5, 5));
         }
-
-        Assert.assertEquals(shapeMap.get("01SequentialBlock"), new Shape(batchSize, 96, 54, 54));
-        Assert.assertEquals(shapeMap.get("03SequentialBlock"), new Shape(batchSize, 256, 26, 26));
-        Assert.assertEquals(shapeMap.get("05SequentialBlock"), new Shape(batchSize, 384, 12, 12));
-        Assert.assertEquals(shapeMap.get("08SequentialBlock"), new Shape(batchSize, 10, 5, 5));
-        manager.close();
     }
 
     @Test
     public void testForwardMethod() {
-        NDManager manager = NDManager.newBaseManager();
-        Block nin = NiN.builder().build();
-        int batchSize = 1;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
-        nin.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        nin.initialize(manager, DataType.FLOAT32, x.getShape());
-        NDArray xHat =
-                nin.forward(new ParameterStore(manager, true), new NDList(x), false)
-                        .singletonOrThrow();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            Block nin = NiN.builder().build();
+            int batchSize = 1;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
+            nin.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            nin.initialize(manager, DataType.FLOAT32, x.getShape());
+            ParameterStore ps = new ParameterStore(manager, true);
+            NDArray xHat = nin.forward(ps, new NDList(x), false).singletonOrThrow();
 
-        Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
-        manager.close();
+            Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
+        }
     }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/VGGTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/VGGTest.java
index 21583996b..ceccffbca 100644
--- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/VGGTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/VGGTest.java
@@ -101,47 +101,51 @@ public class VGGTest {
 
     @Test
     public void testOutputShapes() {
-        NDManager manager = NDManager.newBaseManager();
-        int batchSize = 1;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
-        Shape currentShape = x.getShape();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            int batchSize = 1;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
+            Shape currentShape = x.getShape();
 
-        Block vgg = VGG.builder().build();
-        vgg.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        vgg.initialize(manager, DataType.FLOAT32, currentShape);
+            Block vgg = VGG.builder().build();
+            vgg.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            vgg.initialize(manager, DataType.FLOAT32, currentShape);
 
-        Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
-        for (int i = 0; i < vgg.getChildren().size(); i++) {
+            Map<String, Shape> shapeMap = new ConcurrentHashMap<>();
+            for (int i = 0; i < vgg.getChildren().size(); i++) {
+                Shape[] newShape =
+                        vgg.getChildren()
+                                .get(i)
+                                .getValue()
+                                .getOutputShapes(new Shape[] {currentShape});
+                currentShape = newShape[0];
+                shapeMap.put(vgg.getChildren().get(i).getKey(), currentShape);
+            }
 
-            Shape[] newShape =
-                    vgg.getChildren().get(i).getValue().getOutputShapes(new Shape[] {currentShape});
-            currentShape = newShape[0];
-            shapeMap.put(vgg.getChildren().get(i).getKey(), currentShape);
+            Assert.assertEquals(
+                    shapeMap.get("01SequentialBlock"), new Shape(batchSize, 64, 112, 112));
+            Assert.assertEquals(
+                    shapeMap.get("02SequentialBlock"), new Shape(batchSize, 128, 56, 56));
+            Assert.assertEquals(
+                    shapeMap.get("03SequentialBlock"), new Shape(batchSize, 256, 28, 28));
+            Assert.assertEquals(
+                    shapeMap.get("04SequentialBlock"), new Shape(batchSize, 512, 14, 14));
+            Assert.assertEquals(shapeMap.get("05SequentialBlock"), new Shape(batchSize, 512, 7, 7));
+            Assert.assertEquals(shapeMap.get("07Linear"), new Shape(batchSize, 4096));
         }
-
-        Assert.assertEquals(shapeMap.get("01SequentialBlock"), new Shape(batchSize, 64, 112, 112));
-        Assert.assertEquals(shapeMap.get("02SequentialBlock"), new Shape(batchSize, 128, 56, 56));
-        Assert.assertEquals(shapeMap.get("03SequentialBlock"), new Shape(batchSize, 256, 28, 28));
-        Assert.assertEquals(shapeMap.get("04SequentialBlock"), new Shape(batchSize, 512, 14, 14));
-        Assert.assertEquals(shapeMap.get("05SequentialBlock"), new Shape(batchSize, 512, 7, 7));
-        Assert.assertEquals(shapeMap.get("07Linear"), new Shape(batchSize, 4096));
-        manager.close();
     }
 
     @Test
     public void testForwardMethod() {
-        NDManager manager = NDManager.newBaseManager();
-        Block vgg = VGG.builder().build();
-        int batchSize = 1;
-        NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
-        vgg.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
-        vgg.initialize(manager, DataType.FLOAT32, x.getShape());
-
-        NDArray xHat =
-                vgg.forward(new ParameterStore(manager, true), new NDList(x), false)
-                        .singletonOrThrow();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            Block vgg = VGG.builder().build();
+            int batchSize = 1;
+            NDArray x = manager.ones(new Shape(batchSize, 1, 224, 224));
+            vgg.setInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+            vgg.initialize(manager, DataType.FLOAT32, x.getShape());
+            ParameterStore ps = new ParameterStore(manager, true);
+            NDArray xHat = vgg.forward(ps, new NDList(x), false).singletonOrThrow();
 
-        Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
-        manager.close();
+            Assert.assertEquals(xHat.getShape(), new Shape(batchSize, 10));
+        }
     }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java b/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
index d8ab3958d..9206dcdd4 100644
--- a/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
@@ -658,7 +658,6 @@ public class BlockCoreTest {
     @SuppressWarnings("try")
     @Test
     public void testGRU() throws IOException, MalformedModelException {
-
         Loss loss = new SoftmaxCrossEntropyLoss("SmCeLoss", 1, -1, false, true);
         TrainingConfig config =
                 new DefaultTrainingConfig(loss)
diff --git a/integration/src/main/java/ai/djl/integration/tests/nn/BlockFactoryTest.java b/integration/src/main/java/ai/djl/integration/tests/nn/BlockFactoryTest.java
index 6b79943c9..4f59e0439 100644
--- a/integration/src/main/java/ai/djl/integration/tests/nn/BlockFactoryTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/nn/BlockFactoryTest.java
@@ -76,8 +76,8 @@ public class BlockFactoryTest {
     private Path prepareModel(Path savedDir)
             throws IOException, ModelNotFoundException, MalformedModelException {
         TestBlockFactory factory = new TestBlockFactory();
-        Model model = factory.getRemoveLastBlockModel();
-        try (NDManager manager = NDManager.newBaseManager()) {
+        try (Model model = factory.getRemoveLastBlockModel();
+                NDManager manager = NDManager.newBaseManager()) {
             Block block = model.getBlock();
             block.forward(
                     new ParameterStore(manager, true),
@@ -85,7 +85,6 @@ public class BlockFactoryTest {
                     true);
             model.save(savedDir, "exported");
         }
-        model.close();
         Path classDir = savedDir.resolve("classes/ai/djl/integration/tests/nn");
         Files.createDirectories(classDir);
         Files.copy(
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/EvaluateDatasetTest.java b/integration/src/main/java/ai/djl/integration/tests/training/EvaluateDatasetTest.java
index 0706a9b42..5865a3a33 100644
--- a/integration/src/main/java/ai/djl/integration/tests/training/EvaluateDatasetTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/training/EvaluateDatasetTest.java
@@ -49,33 +49,33 @@ public class EvaluateDatasetTest {
 
             Mlp mlpModel = new Mlp(784, 1, new int[] {256}, Activation::relu);
 
-            Model model = Model.newInstance("lin-reg");
-            model.setBlock(mlpModel);
+            try (Model model = Model.newInstance("lin-reg")) {
+                model.setBlock(mlpModel);
 
-            Loss l2loss = Loss.l2Loss();
+                Loss l2loss = Loss.l2Loss();
 
-            Tracker lrt = Tracker.fixed(0.5f);
-            Optimizer sgd = Optimizer.sgd().setLearningRateTracker(lrt).build();
+                Tracker lrt = Tracker.fixed(0.5f);
+                Optimizer sgd = Optimizer.sgd().setLearningRateTracker(lrt).build();
 
-            DefaultTrainingConfig config =
-                    new DefaultTrainingConfig(l2loss)
-                            .optOptimizer(sgd) // Optimizer (loss function)
-                            .addTrainingListeners(TrainingListener.Defaults.logging()); // Logging
+                DefaultTrainingConfig config =
+                        new DefaultTrainingConfig(l2loss)
+                                .optOptimizer(sgd) // Optimizer (loss function)
+                                .addTrainingListeners(
+                                        TrainingListener.Defaults.logging()); // Logging
 
-            try (Trainer trainer = model.newTrainer(config)) {
-                trainer.initialize(new Shape(1, 784));
+                try (Trainer trainer = model.newTrainer(config)) {
+                    trainer.initialize(new Shape(1, 784));
 
-                Metrics metrics = new Metrics();
-                trainer.setMetrics(metrics);
+                    Metrics metrics = new Metrics();
+                    trainer.setMetrics(metrics);
 
-                EasyTrain.evaluateDataset(trainer, testMnistDataset);
+                    EasyTrain.evaluateDataset(trainer, testMnistDataset);
 
-                Assert.assertTrue(
-                        l2loss.getAccumulator(EvaluatorTrainingListener.VALIDATE_EPOCH) < 25.0f,
-                        "dataset L2 loss is more than expected.");
+                    Assert.assertTrue(
+                            l2loss.getAccumulator(EvaluatorTrainingListener.VALIDATE_EPOCH) < 25.0f,
+                            "dataset L2 loss is more than expected.");
+                }
             }
-
-            model.close();
         }
     }
 }