try datatype

Author: pcmoritz

cpp/src/plasma/tensorflow_op/tf_plasma_op.cc

@@ -284,7 +284,7 @@ class TensorToPlasmaOp : public AsyncOpKernel {
       for (int i = 0; i < num_tensors; ++i) {
         const auto& input_tensor = context->input(i);
         std::memcpy(static_cast<void*>(data + offsets[i] / byte_width),
-                    input_tensor.flat<float>().data(),
+                    input_tensor.flat<tf_dtype::Type>().data(),
                     static_cast<uint64>(offsets[i + 1] - offsets[i]));
       }
       wrapped_callback();