Fix memory leak in zeroGradients().

Author: neutrino

engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java

@@ -123,7 +123,10 @@ public void zeroGradients() {
         NDManager systemManager = MxNDManager.getSystemManager();
         for (NDArray array : systemManager.getManagedArrays()) {
             if (array.hasGradient()) {
-                array.getGradient().subi(array.getGradient());
+                // To prevent memory leak we must close gradient after use.
+                try (NDArray gradient = array.getGradient()) {
+                    gradient.subi(gradient);
+                }
             }
         }
     }

engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java

@@ -75,8 +75,9 @@ private void backward(NDArray target, NDArray grad, boolean keepGraph, boolean c
     public void zeroGradients() {
         NDManager systemManager = PtNDManager.getSystemManager();
         for (NDArray array : systemManager.getManagedArrays()) {
-            if (array.hasGradient()) {
-                array.getGradient().subi(array.getGradient());
+            // To prevent memory leak we must close gradient after use.
+            try (NDArray gradient = array.getGradient()) {
+                gradient.subi(gradient);
             }
         }
     }