[ROCM] fix reduce_sum nan in ROCM platform, test=develop (#31780)

4 years ago · 46dd1d4aad
parent f72d197ec5
commit 46dd1d4aad
1 changed files with 4 additions and 0 deletions
--- a/paddle/fluid/operators/reduce_ops/cub_reduce.h
+++ b/paddle/fluid/operators/reduce_ops/cub_reduce.h
@ -161,7 +161,11 @@ static inline std::vector<int> GetStrides(const std::vector<int>& dims,
  return strides;
 }

+#ifdef __HIPCC__
+constexpr int kMaxBlockDim = 256;
+#else
 constexpr int kMaxBlockDim = 512;
+#endif

 static inline int GetDesiredBlockDim(int block_dim) {
  return block_dim >= kMaxBlockDim