diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2017-12-04 00:57:14 -0800 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2017-12-04 01:00:36 -0800 |
| commit | 88e63eb7099e2cd82e942f4b3867a9fedffb3d85 (patch) | |
| tree | 7dc802f1e1af948f15639465aca4212130dea7e1 /tensorflow/core/kernels/scatter_nd_op_gpu.cu.cc | |
| parent | ed2e8c227e77d591b868ecd6d1c8d59ddae3b3d9 (diff) | |
Make gather_nd & scatter_nd handle rank 6 and 7 indices.
PiperOrigin-RevId: 177781215
Diffstat (limited to 'tensorflow/core/kernels/scatter_nd_op_gpu.cu.cc')
| -rw-r--r-- | tensorflow/core/kernels/scatter_nd_op_gpu.cu.cc | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/tensorflow/core/kernels/scatter_nd_op_gpu.cu.cc b/tensorflow/core/kernels/scatter_nd_op_gpu.cu.cc index 0eb3cf32dd..31f74671ca 100644 --- a/tensorflow/core/kernels/scatter_nd_op_gpu.cu.cc +++ b/tensorflow/core/kernels/scatter_nd_op_gpu.cu.cc @@ -136,7 +136,9 @@ struct ScatterNdFunctor<GPUDevice, T, Index, op, IXDIM> { DECLARE_GPU_SPECS_INDEX_OP_IXDIM(T, Index, op, 2); \ DECLARE_GPU_SPECS_INDEX_OP_IXDIM(T, Index, op, 3); \ DECLARE_GPU_SPECS_INDEX_OP_IXDIM(T, Index, op, 4); \ - DECLARE_GPU_SPECS_INDEX_OP_IXDIM(T, Index, op, 5); + DECLARE_GPU_SPECS_INDEX_OP_IXDIM(T, Index, op, 5); \ + DECLARE_GPU_SPECS_INDEX_OP_IXDIM(T, Index, op, 6); \ + DECLARE_GPU_SPECS_INDEX_OP_IXDIM(T, Index, op, 7); #define DECLARE_GPU_SPECS_INDEX(T, Index) \ DECLARE_GPU_SPECS_INDEX_OP(T, Index, scatter_nd_op::UpdateOp::ASSIGN); \ |