Add amend

apache · Jul 22, 2018 · 203d237 · 203d237
1 parent 49f37a9
commit 203d237
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/io/MXDataIter.scala b/scala-package/core/src/main/scala/org/apache/mxnet/io/MXDataIter.scala
@@ -125,7 +125,8 @@ private[mxnet] class MXDataIter(private[mxnet] val handle: DataIterHandle,
     checkCall(_LIB.mxDataIterNext(handle, next))
     if (next.value > 0) {
       currentBatch = new DataBatch(data = getData(), label = getLabel(),
-        index = getIndex(), pad = getPad())
+        index = getIndex(), pad = getPad(),
+        dtype = currentBatch.dtype, layout = currentBatch.layout)
     } else {
       currentBatch = null
     }