fix shard_index kernel (#46491)

PaddlePaddle · Sep 26, 2022 · 808bf2b · 808bf2b
1 parent d150c3a
commit 808bf2b
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/paddle/phi/kernels/gpu/shard_index_kernel.cu b/paddle/phi/kernels/gpu/shard_index_kernel.cu
@@ -33,7 +33,15 @@ __global__ void ShardIndexInner(const T* in_data,
   int shard_size = (index_num + nshards - 1) / nshards;
   int idx = blockIdx.x * blockDim.x + threadIdx.x;
   if (idx < numel) {
-    assert(in_data[idx] >= 0 && in_data[idx] < index_num);
+    PADDLE_ENFORCE(in_data[idx] >= 0,
+                   "The input_index for Op(shard_index) must be "
+                   "greater or equal to 0, but the value given is %d.",
+                   in_data[idx]);
+    PADDLE_ENFORCE(in_data[idx] < index_num,
+                   "The input_index for Op(shard_index) must be less "
+                   "than index_num (%d), but the value given is %d.",
+                   index_num,
+                   in_data[idx]);
     if (in_data[idx] / shard_size == shard_id) {
       out_data[idx] = in_data[idx] % shard_size;
     } else {