Fix memory leak in zeroGradients().

deepjavalibrary · Sep 29, 2023 · 7d61e5c · 7d61e5c
1 parent 963332d
commit 7d61e5c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java b/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);
+                }
             }
         }
     }

diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java b/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);
             }
         }
     }