diff options
author | 2016-06-06 12:45:58 -0800 | |
---|---|---|
committer | 2016-06-06 14:04:47 -0700 | |
commit | aba8beebab0b363f03492b3d5653ec14d148f3c3 (patch) | |
tree | 82648ff178ba1bf4fc5232328f2bc650cde27da6 /tensorflow/core/kernels/avgpooling_op.cc | |
parent | e70c452a18403b368ea845cfb654079386a00fd8 (diff) |
Change some kernels to use TF_CALL* macros, so that the instantiations for some
types can be avoided on mobile platform.
Change: 124172890
Diffstat (limited to 'tensorflow/core/kernels/avgpooling_op.cc')
-rw-r--r-- | tensorflow/core/kernels/avgpooling_op.cc | 20 |
1 files changed, 10 insertions, 10 deletions
diff --git a/tensorflow/core/kernels/avgpooling_op.cc b/tensorflow/core/kernels/avgpooling_op.cc index fc7f6d1a5a..4378dd2fa4 100644 --- a/tensorflow/core/kernels/avgpooling_op.cc +++ b/tensorflow/core/kernels/avgpooling_op.cc @@ -23,6 +23,7 @@ limitations under the License. #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/framework/numeric_op.h" #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_slice.h" @@ -337,16 +338,15 @@ class AvgPoolingGradOp : public OpKernel { TensorFormat data_format_; }; -REGISTER_KERNEL_BUILDER(Name("AvgPoolGrad") - .Device(DEVICE_CPU) - .TypeConstraint<float>("T") - .HostMemory("orig_input_shape"), - AvgPoolingGradOp<CPUDevice, float>); -REGISTER_KERNEL_BUILDER(Name("AvgPoolGrad") - .Device(DEVICE_CPU) - .TypeConstraint<double>("T") - .HostMemory("orig_input_shape"), - AvgPoolingGradOp<CPUDevice, double>); +#define REGISTER_CPU_KERNEL(T) \ + REGISTER_KERNEL_BUILDER(Name("AvgPoolGrad") \ + .Device(DEVICE_CPU) \ + .TypeConstraint<T>("T") \ + .HostMemory("orig_input_shape"), \ + AvgPoolingGradOp<CPUDevice, T>); + +TF_CALL_float(REGISTER_CPU_KERNEL); +TF_CALL_double(REGISTER_CPU_KERNEL); #if GOOGLE_CUDA |