diff --git a/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/Client.scala b/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
index 15f3c4f180ea3..9504b99de3246 100644
--- a/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
+++ b/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
@@ -113,11 +113,12 @@ class Client(clientArgs: ClientArguments, hadoopConf: Configuration, spConf: Spa
   }
 
   def calculateAMMemory(newApp: GetNewApplicationResponse) :Int = {
+    // It looks like -Xmx is handled in ClientBase. ?
     // TODO: Need a replacement for the following code to fix -Xmx?
-    // val minResMemory: Int = newApp.getMinimumResourceCapability().getMemory()
-    // var amMemory = ((args.amMemory / minResMemory) * minResMemory) +
-    //  ((if ((args.amMemory % minResMemory) == 0) 0 else minResMemory) -
-    //    memoryOverhead )
+    val minResMemory: Int = newApp.getMinimumResourceCapability().getMemory()
+    var amMemory = ((args.amMemory / minResMemory) * minResMemory) +
+     ((if ((args.amMemory % minResMemory) == 0) 0 else minResMemory) -
+       memoryOverhead )
     args.amMemory
   }