diff --git a/paddle/phi/kernels/gpu/multinomial_kernel.cu b/paddle/phi/kernels/gpu/multinomial_kernel.cu index ef6cd1323a9df832e0051d78d522cab8eb00121c..21a506a840cc7f315f7703bc950d0e6c1fba88e5 100644 --- a/paddle/phi/kernels/gpu/multinomial_kernel.cu +++ b/paddle/phi/kernels/gpu/multinomial_kernel.cu @@ -133,11 +133,10 @@ void MultinomialKernel(const Context& dev_ctx, DenseTensor* out) { auto* in_data = x.data(); int64_t* out_data = dev_ctx.template Alloc(out); - auto in_dims = x.dims(); - int64_t in_rank = in_dims.size(); - const int64_t num_categories = in_dims[in_rank - 1]; - const int64_t num_distributions = in_rank > 1 ? in_dims[in_rank - 2] : 1; + int64_t dim_size = in_dims.size(); + const int64_t num_categories = in_dims[dim_size - 1]; + const int64_t num_distributions = dim_size > 1 ? in_dims[dim_size - 2] : 1; // If replacement is False, it's not a replaceable sample. Every category // can be used only once. @@ -145,8 +144,8 @@ void MultinomialKernel(const Context& dev_ctx, int64_t in_data_numel = x.numel(); int64_t out_data_numel = out->numel(); + // Just use to PADDLE_ENFORCE error message T* cpu_in_data = new T[in_data_numel]; - int64_t* cpu_out_data = new int64_t[out_data_numel]; #ifdef PADDLE_WITH_HIP hipMemcpy( @@ -160,7 +159,7 @@ void MultinomialKernel(const Context& dev_ctx, for (size_t i = 0; i < num_distributions; ++i) { int zero_num = 0; for (size_t j = 0; j < num_categories; ++j) { - T weight = cpu_in_data[i * num_distributions + j]; + T weight = cpu_in_data[i * num_categories + j]; PADDLE_ENFORCE_GE( weight, 0, diff --git a/python/paddle/fluid/tests/unittests/test_multinomial_op.py b/python/paddle/fluid/tests/unittests/test_multinomial_op.py index ecde527523d3dd45eecdbb5aeba468c0456bc5b0..4dfc881d7723f919ad61cf67086e945f9f8ee6cb 100644 --- a/python/paddle/fluid/tests/unittests/test_multinomial_op.py +++ b/python/paddle/fluid/tests/unittests/test_multinomial_op.py @@ -216,6 +216,14 @@ class TestMultinomialError(unittest.TestCase): self.assertRaises(ValueError, test_dim_less_than_1) + with self.assertRaises(ValueError): + y = paddle.multinomial(paddle.to_tensor([1., 2., -3.])) + + with self.assertRaises(ValueError): + prob = paddle.rand([20, 1000]) + prob[1:0] = 0 + y = paddle.multinomial(prob) + class TestRandomValue(unittest.TestCase): def test_fixed_random_number(self):