OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 12:20:11 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Merge pull request #50700 from geetachavan1/cherrypicks_FSBQ7

Browse Source 
[CherryPick 2.6]Disable GPU kernel for SparseSegmentMeanGrad
This commit is contained in:
Mihai Maruseac 2021-07-09 10:31:00 -07:00 committed by GitHub
commit 4a31b52c6c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/core/kernels/segment_reduction_ops_impl_5.cc
View File

@ -217,6 +217,7 @@ TF_CALL_FLOAT_TYPES(REGISTER_CPU_SPARSE_KERNELS_FOR_EACH_INDEX_TYPE);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_SPARSE_KERNELS_FOR_EACH_INDEX_TYPE);
#undef REGISTER_GPU_SPARSE_KERNELS
#if 0 // TODO(b/192086735): Enable once bug is fixed.
#define REGISTER_GPU_SPARSE_KERNELS(type, index_type, segment_ids_type) \
REGISTER_KERNEL_BUILDER( \
Name("SparseSegmentMeanGrad") \
@ -228,6 +229,7 @@ TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_SPARSE_KERNELS_FOR_EACH_INDEX_TYPE);
SparseSegmentMeanGradOp<GPUDevice, type, index_type, segment_ids_type>);
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_SPARSE_KERNELS_FOR_EACH_INDEX_TYPE);
#undef REGISTER_GPU_SPARSE_KERNELS
#endif
#define REGISTER_GPU_SPARSE_KERNELS(type, index_type, segment_ids_type) \
REGISTER_KERNEL_BUILDER( \