add check on input probability (#59174)

zhiqiu · web-flow · commit 16aa5c4cb8ed · 2023-11-21T12:01:07.000+08:00
diff --git a/paddle/phi/kernels/gpu/bernoulli_kernel.cu b/paddle/phi/kernels/gpu/bernoulli_kernel.cu
@@ -57,8 +57,11 @@ __global__ void bernoulli_cuda_kernel(
     for (size_t j = 0; j < 4; j++) {
       size_t idx = i + j;
       if (idx < size) {
-        out_data[idx] =
-            static_cast<T>((&rand.x)[j] <= static_cast<MPType>(x_data[idx]));
+        MPType p = static_cast<MPType>(x_data[idx]);
+        PADDLE_ENFORCE(p >= 0 && p <= 1,
+                       "The probability should be in [0, 1], but got %f",
+                       p);
+        out_data[idx] = static_cast<T>((&rand.x)[j] <= static_cast<MPType>(p));
       }
     }
   }

Original file line number	Diff line number	Diff line change
`@@ -57,8 +57,11 @@ __global__ void bernoulli_cuda_kernel(`
`57`	`57`	`for (size_t j = 0; j < 4; j++) {`
`58`	`58`	`size_t idx = i + j;`
`59`	`59`	`if (idx < size) {`
`60`		`- out_data[idx] =`
`61`		`- static_cast<T>((&rand.x)[j] <= static_cast<MPType>(x_data[idx]));`
	`60`	`+ MPType p = static_cast<MPType>(x_data[idx]);`
	`61`	`+ PADDLE_ENFORCE(p >= 0 && p <= 1,`
	`62`	`+ "The probability should be in [0, 1], but got %f",`
	`63`	`+ p);`
	`64`	`+ out_data[idx] = static_cast<T>((&rand.x)[j] <= static_cast<MPType>(p));`
`62`	`65`	`}`
`63`	`66`	`}`
`64`	`67`	`}`