[api] Fixes IdEmbedding memory leak (#3257)

Fixes: #3251
deepjavalibrary · Jun 17, 2024 · bb7dfd1 · bb7dfd1
1 parent 34caf33
commit bb7dfd1
Showing 1 changed file with 14 additions and 9 deletions.
diff --git a/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java b/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java
@@ -63,15 +63,20 @@ public Shape[] getOutputShapes(Shape[] inputShapes) {
     protected NDList forwardInternal(
             ParameterStore ps, NDList inputs, boolean training, PairList<String, Object> params) {
         NDArray input = inputs.singletonOrThrow();
-        // on info to the right shapes, see: http://beta.mxnet.io/r/api/mx.symbol.gather_nd.html
-        NDArray ids = input.flatten().reshape(1, input.getShape().size());
-        // create the embedding Table
-        NDArray embeddingTable = ps.getValue(embedding, ids.getDevice(), training);
-        // We do not perform a sparse lookup, instead we just project into the table
-        NDArray result = embeddingTable.gatherNd(ids);
-        // we want the original shape of the input + the last dimension of the embedding
-        Shape targetShape = input.getShape().addAll(new Shape(embeddingTable.getShape().get(1)));
-        return new NDList(result.reshape(targetShape));
+        try (NDManager scope = NDManager.subManagerOf(input)) {
+            // on info to the right shapes, see: http://beta.mxnet.io/r/api/mx.symbol.gather_nd.html
+            NDArray ids = input.flatten().reshape(1, input.getShape().size());
+            // create the embedding Table
+            NDArray embeddingTable = ps.getValue(embedding, ids.getDevice(), training);
+            scope.tempAttachAll(embeddingTable);
+            // We do not perform a sparse lookup, instead we just project into the table
+            NDArray result = embeddingTable.gatherNd(ids);
+            result.attach(inputs.getManager());
+            // we want the original shape of the input + the last dimension of the embedding
+            Shape targetShape =
+                    input.getShape().addAll(new Shape(embeddingTable.getShape().get(1)));
+            return new NDList(result.reshape(targetShape));
+        }
     }
 
     /**