// Simple MLP network
def mlpNetwork(): Symbol = {
    val input = Symbol.Variable("data")
    val label = Symbol.Variable("label")
    val fc1 = Symbol.FullyConnected(name = "fc1")()(Map("data" -> input, "num_hidden" -> 128))
    val act1 = Symbol.Activation(name = "relu")()(Map("data" -> fc1, "act_type" -> "relu"))
    val fc2 = Symbol.FullyConnected(name = "fc2")()(Map("data" -> act1, "num_hidden" -> 1))
    val loss = Symbol.LinearRegressionOutput(name="loss")()(Map("data" -> fc2, "label" -> label))
    loss
  }

def getNDArrayIter: NDArrayIter = {
    val f = NDArray.zeros(100, 20, 20)
    val l = NDArray.zeros(100, 1)
    val data = Array(f)
    val labels = Array(l)
    val batchSize = 10
    val iter = new NDArrayIter(data, labels, batchSize)
    iter
  }

val net = mlpNetwork()
val iter = getNDArrayIter()
val optimizer = new Adam(0.001f, 0.9f, 0.999f, 1e-8f, 1 - 1e-8f, 0f, 10f, null);
val init = new Normal(0.01f);
val model = FeedForward.newBuilder(modelSpec)
                .setContext(Array(Context.gpu(0)))
                .setInitializer(init)
                .setNumEpoch(100000)
                .setOptimizer(optimizer)
                .setTrainData(iter)
                .setEvalData(iter)
                .build();

val meanT = (beta1t * mean + (1.0 - beta1t) * resdGrad)
      .disposeDepsExcept(mean, resdGrad)

val varianceT = (beta2 * variance + (1.0f - beta2) * resdGrad * resdGrad)
      .disposeDepsExcept(variance, resdGrad)

val step = (learningRate * meanT / (NDArray.sqrt(varianceT) + epsilon))
      .disposeDepsExcept(meanT, varianceT)

val beta1Mean = beta1 * mean
    val beta1ResGrad = (1.0 - beta1t) * resdGrad
    val meanT = beta1Mean + beta1ResGrad
    // dipose temp NDArrays
    betaMean.dispose()
    betaResGrad.dispose()

val beta2Variance = beta2 * variance
    val beta2ResGrad = (1.0f - beta2) * resdGrad
    val beta2ResGradSquare = beta2ResGrad * resdGrad
    val varianceT = beta2Variance + beta2ResGradSquare
    // dipose temp NDArrays
    beta2Variance.dispose()
    beta2ResGrad.dispose()
    beta2ResGradSquare.dispose()


    val lrMeanT = learningRate * meanT
    val sqrtVar = NDArray.sqrt(varianceT)
    val sqrtVarPlusEpsilon = sqrtVar + epsilon
    val step = lrMeanT / sqrtVarPlusEpsilon
    // dipose temp NDArrays
    lrMeanT.dispose()
    sqrtVar.dispose()
    sqrtVarPlusEpsilon.dispose()