experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Grappler] Register DataFormatDimMap on host for CPU.

Browse Source 
PiperOrigin-RevId: 256262935
This commit is contained in:
Andy Ly 2019-07-02 16:29:35 -07:00 committed by TensorFlower Gardener
parent 7f1b6e06ba
commit 237ce08c89
1 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tensorflow/core/kernels/data_format_ops.cc
View File

@ -156,6 +156,16 @@ TF_CALL_int32(REGISTER_KERNEL);
TF_CALL_int64(REGISTER_KERNEL);
#undef REGISTER_KERNEL
#define REGISTER_KERNEL(T) \
REGISTER_KERNEL_BUILDER(Name("DataFormatDimMap") \
.Device(DEVICE_CPU) \
.Label("host") \
.TypeConstraint<T>("T"), \
DataFormatDimMapOp<CPUDevice, T>);
TF_CALL_int32(REGISTER_KERNEL);
TF_CALL_int64(REGISTER_KERNEL);
#undef REGISTER_KERNEL
#define REGISTER_KERNEL(T) \
REGISTER_KERNEL_BUILDER(Name("DataFormatVecPermute") \
.Device(DEVICE_CPU) \