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

Adding ROCm support for batch_to_space op

Browse Source
This commit is contained in:
Deven Desai 2019-07-03 03:07:04 +00:00
parent 28d194b368
commit 87ec8b8827
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/core/kernels/batchtospace_op.cc
View File

@ -266,7 +266,7 @@ class BatchToSpaceOp : public OpKernel {
TF_CALL_REAL_NUMBER_TYPES(REGISTER);
#undef REGISTER
#if GOOGLE_CUDA
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#define REGISTER(T) \
REGISTER_KERNEL_BUILDER(Name("BatchToSpaceND") \
.Device(DEVICE_GPU) \
@ -282,6 +282,6 @@ TF_CALL_REAL_NUMBER_TYPES(REGISTER);
TF_CALL_GPU_NUMBER_TYPES(REGISTER);
#undef REGISTER
#endif // GOOGLE_CUDA
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} // end namespace tensorflow