Require min/max operands to be on the CPU.

Change: 150028654

Require min/max operands to be on the CPU.
Change: 150028654
5388b228 · Patrick Nguyen · TensorFlower Gardener · 3971c27b · 5388b228
隐藏空白更改
内联并排

Showing with 2 addition and 0 deletion

tensorflow/core/kernels/quantize_and_dequantize_op.cc tensorflow/core/kernels/quantize_and_dequantize_op.cc +2 -0

未找到文件。
--- a/tensorflow/core/kernels/quantize_and_dequantize_op.cc
+++ b/tensorflow/core/kernels/quantize_and_dequantize_op.cc
@@ -164,6 +164,8 @@ TF_CALL_double(REGISTER_CPU_KERNEL);
 #define REGISTER_GPU_KERNEL(T)                                                 \
  REGISTER_KERNEL_BUILDER(Name("QuantizeAndDequantizeV2")                      \
                              .Device(DEVICE_GPU)                              \
+                              .HostMemory("input_max")                         \
+                              .HostMemory("input_min")                         \
                              .TypeConstraint<T>("T"),                         \
                          QuantizeAndDequantizeV2Op<GPUDevice, T>);            \
  REGISTER_KERNEL_BUILDER(                                                     \