[ROCM] fix softmax with loss nan in HIP platform, test=develop

PaddlePaddle · Mar 9, 2021 · be0a056 · be0a056 · paddle-bot-old · Mar 9, 2021
1 parent 133a914
commit be0a056
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/paddle/fluid/operators/softmax_with_cross_entropy_op.cu b/paddle/fluid/operators/softmax_with_cross_entropy_op.cu
@@ -398,7 +398,12 @@ static void HardLabelSoftmaxWithCrossEntropy(
     const platform::CUDADeviceContext& ctx, const T* logits_data,
     const int64_t* labels_data, T* loss_data, T* softmax_data, int64_t n,
     int64_t d, int axis_dim, int ignore_idx) {
+#ifdef __HIPCC__
+  // HIP platform will have loss nan if dim size > 256
+  constexpr int kMaxBlockDim = 256;
+#else
   constexpr int kMaxBlockDim = 512;
+#endif
   int64_t block_dim = axis_dim >= kMaxBlockDim
                           ? kMaxBlockDim
                           : (1 << static_cast<int>(std::log2(axis_dim)));