Support bool value in tf dataset (intel#2401)

* support bool value in tf dataset * fix style
dding3 · Jun 3, 2020 · 2739898 · 2739898
1 parent 16e7605
commit 2739898
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/common/TFUtils.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/common/TFUtils.scala
@@ -96,7 +96,7 @@ object TFUtils {
     val dataType = t.dataType()
 
     val numericDataTypes = Set(DataType.FLOAT,
-      DataType.UINT8, DataType.INT32, DataType.INT64, DataType.DOUBLE)
+      DataType.UINT8, DataType.INT32, DataType.INT64, DataType.DOUBLE, DataType.BOOL)
 
     if (dataType == DataType.STRING) {
       val outputTensor = output.asInstanceOf[Tensor[Array[Byte]]]
@@ -148,6 +148,13 @@ object TFUtils {
           val buffer = DoubleBuffer.wrap(arr)
           t.writeTo(buffer)
           double2float(arr, outputTensor.storage().array(), outputTensor.storageOffset() - 1)
+        case DataType.BOOL =>
+          val outputTensor = output.asInstanceOf[Tensor[Float]]
+          val arr = new Array[Byte](t.numBytes())
+          assert(t.numBytes() == shape.product, "sanity check")
+          val buffer = ByteBuffer.wrap(arr)
+          t.writeTo(buffer)
+          byte2float(arr, outputTensor.storage().array(), outputTensor.storageOffset() - 1)
       }
 
     } else {