From a5ddc2fa5916522d18fc2ee33d22ecd1843212e1 Mon Sep 17 00:00:00 2001
From: zhijxu <zhijxu@microsoft.com>
Date: Tue, 26 Mar 2024 08:43:37 +0800
Subject: [PATCH] resolve review comments

---
 .../core/providers/cuda/math/softmax_warpwise_impl.cuh       | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/onnxruntime/core/providers/cuda/math/softmax_warpwise_impl.cuh b/onnxruntime/core/providers/cuda/math/softmax_warpwise_impl.cuh
index 054f9abc8c071..5e2cec464a86b 100644
--- a/onnxruntime/core/providers/cuda/math/softmax_warpwise_impl.cuh
+++ b/onnxruntime/core/providers/cuda/math/softmax_warpwise_impl.cuh
@@ -184,8 +184,9 @@ __global__ void softmax_warp_forward_resource_efficient(output_t* dst, const inp
     if (element_index < element_count) {
       elements[it][local_idx] = src[it * WARP_SIZE];
     } else {
-      static_assert(!std::is_same<acc_t, half>::value, "acc_t can no be half, as the infinity function will return 0 instead of inf");
-      elements[it][local_idx] = (input_t)-std::numeric_limits<acc_t>::infinity();
+      static_assert(std::numeric_limits<acc_t>::has_infinity,
+                    "type of acc_t should have infinity to avoid infinity function return 0");
+      elements[it][local_idx] = static_cast<input_t>(-std::numeric_limits<acc_t>::infinity());
     }
   }
   // compute max_value