Gan and Customop fixes

apache · Jul 24, 2018 · 7340653 · 7340653
1 parent 836e64f
commit 7340653
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/scala-package/examples/src/main/scala/org/apache/mxnetexamples/customop/Data.scala b/scala-package/examples/src/main/scala/org/apache/mxnetexamples/customop/Data.scala
@@ -20,6 +20,7 @@ package org.apache.mxnetexamples.customop
 import org.apache.mxnet.{DataIter, IO, Shape}
 
 object Data {
+
   // return train and val iterators for mnist
   def mnistIterator(dataPath: String, batchSize: Int, inputShape: Shape): (DataIter, DataIter) = {
     val flat = if (inputShape.length == 3) "False" else "True"
@@ -29,15 +30,19 @@ object Data {
       "input_shape" -> inputShape.toString(),
       "batch_size" -> s"$batchSize",
       "shuffle" -> "True",
-      "flat" -> flat
+      "flat" -> flat,
+      "dataLayout" -> "NT",
+      "labelLayout" -> "N"
     )
     val trainDataIter = IO.MNISTIter(trainParams)
     val testParams = Map(
       "image" -> s"$dataPath/t10k-images-idx3-ubyte",
       "label" -> s"$dataPath/t10k-labels-idx1-ubyte",
       "input_shape" -> inputShape.toString(),
       "batch_size" -> s"$batchSize",
-      "flat" -> flat
+      "flat" -> flat,
+      "dataLayout" -> "NT",
+      "labelLayout" -> "N"
     )
     val testDataIter = IO.MNISTIter(testParams)
     (trainDataIter, testDataIter)

diff --git a/scala-package/examples/src/main/scala/org/apache/mxnetexamples/gan/GanMnist.scala b/scala-package/examples/src/main/scala/org/apache/mxnetexamples/gan/GanMnist.scala
@@ -130,7 +130,9 @@ object GanMnist {
       "label" -> s"$dataPath/train-labels-idx1-ubyte",
       "input_shape" -> s"(1, 28, 28)",
       "batch_size" -> s"$batchSize",
-      "shuffle" -> "True"
+      "shuffle" -> "True",
+      "dataLayout" -> "NT",
+      "labelLayout" -> "N"
     )
 
     val mnistIter = IO.MNISTIter(params)