Repository: deepjavalibrary/djl
Pull request #1255 (2021-09-27 20:58:49)

Fix issue when use file name to load model.
Change-Id: I7e360dc9899e047444eb93c0af9605495b0882a9

## 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



Patch:
diff --git a/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java b/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
index 95113f791..39f16d0cd 100644
--- a/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
+++ b/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
@@ -98,8 +98,12 @@ public class BaseModelLoader implements ModelLoader {
             }
 
             Path modelPath = mrl.getRepository().getResourceDirectory(artifact);
+            Path modelDir = Files.isRegularFile(modelPath) ? modelPath.getParent() : modelPath;
+            if (modelDir == null) {
+                throw new AssertionError("Directory should not be null.");
+            }
 
-            loadServingProperties(modelPath, arguments, options);
+            loadServingProperties(modelDir, arguments, options);
             Application application = criteria.getApplication();
             if (application != Application.UNDEFINED) {
                 arguments.put("application", application.getPath());
@@ -143,7 +147,7 @@ public class BaseModelLoader implements ModelLoader {
 
             Model model =
                     createModel(
-                            modelPath,
+                            modelDir,
                             modelName,
                             criteria.getDevice(),
                             criteria.getBlock(),
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();
         }
     }
 }