File tree Expand file tree Collapse file tree 2 files changed +8
-2
lines changed
mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine
pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine Expand file tree Collapse file tree 2 files changed +8
-2
lines changed Original file line number Diff line number Diff line change @@ -123,7 +123,10 @@ public void zeroGradients() {
123
123
NDManager systemManager = MxNDManager .getSystemManager ();
124
124
for (NDArray array : systemManager .getManagedArrays ()) {
125
125
if (array .hasGradient ()) {
126
- array .getGradient ().subi (array .getGradient ());
126
+ // To prevent memory leak we must close gradient after use.
127
+ try (NDArray gradient = array .getGradient ()) {
128
+ gradient .subi (gradient );
129
+ }
127
130
}
128
131
}
129
132
}
Original file line number Diff line number Diff line change @@ -76,7 +76,10 @@ public void zeroGradients() {
76
76
NDManager systemManager = PtNDManager .getSystemManager ();
77
77
for (NDArray array : systemManager .getManagedArrays ()) {
78
78
if (array .hasGradient ()) {
79
- array .getGradient ().subi (array .getGradient ());
79
+ // To prevent memory leak we must close gradient after use.
80
+ try (NDArray gradient = array .getGradient ()) {
81
+ gradient .subi (gradient );
82
+ }
80
83
}
81
84
}
82
85
}
You can’t perform that action at this time.
0 commit comments