From 014338aa64ccaca0a179145676eba5a76103e81a Mon Sep 17 00:00:00 2001
From: Andrew Ayres <andrew.f.ayres@gmail.com>
Date: Thu, 14 Jun 2018 10:50:43 -0700
Subject: [PATCH] Scala inference memory leak fix #11204 (#11216)

* Fixes Scala memory leak (#10436)

* Replaced the copy and disposed of sliced ndArray to resolve memory leak

* Wrapped disposes in a finally to ensure they are called.
---
 .../scala/org/apache/mxnet/FeedForward.scala  | 25 +++++++++++++------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/FeedForward.scala b/scala-package/core/src/main/scala/org/apache/mxnet/FeedForward.scala
index 7289df197125..87c9bc72be0c 100644
--- a/scala-package/core/src/main/scala/org/apache/mxnet/FeedForward.scala
+++ b/scala-package/core/src/main/scala/org/apache/mxnet/FeedForward.scala
@@ -224,13 +224,24 @@ class FeedForward private(
     var i = 0
     while (data.hasNext && i != numBatch) {
       val batch = data.next()
-      i += 1
-      ExecutorManager.loadData(batch, dataArrays)
-      predExec.forward(isTrain = false)
-      val padded = batch.pad
-      val realSize = batchSize - padded
-      for ((list, nd) <- outputs zip predExec.outputs) {
-        list += nd.slice(0, realSize).copy()
+      try {
+        i += 1
+        ExecutorManager.loadData(batch, dataArrays)
+        predExec.forward(isTrain = false)
+        val padded = batch.pad
+        val realSize = batchSize - padded
+        for ((list, nd) <- outputs zip predExec.outputs) {
+          // The slice is being written to a value so that dispose can be called after the copy.
+          // The one liner nd.slice().copy() leads to leaking the memory of the slice.
+          val ndSliced = nd.slice(0, realSize)
+          try {
+            list += ndSliced.copy()
+          } finally {
+            ndSliced.dispose()
+          }
+        }
+      } finally {
+        batch.dispose()
       }
     }
     // TODO(Yizhi): we can use Symbol.concat to do the same thing. Can it be more efficient?