未验证 提交 29453da1 编写于 作者: Z Zhang Zheng 提交者: GitHub

Fix bug when eigen_device() is nullptr in top_k (#40459)

上级 88ec08a7
...@@ -78,15 +78,16 @@ void TopkKernel(const Context& dev_ctx, ...@@ -78,15 +78,16 @@ void TopkKernel(const Context& dev_ctx,
// The conclusion is drawn from the data through multiple sets of // The conclusion is drawn from the data through multiple sets of
// statistics // statistics
if (input_width >= 128 && k >= input_width * 0.75) { if (input_width >= 128 && k >= input_width * 0.75) {
if (ops::SortTopk<T>( auto* ctx = reinterpret_cast<const paddle::platform::CUDADeviceContext*>(
paddle::platform::CUDADeviceContext(dev_ctx.GetPlace()), &dev_ctx);
input, if (ops::SortTopk<T>(*ctx,
input_width, input,
input_height, input_width,
k, input_height,
out, k,
indices, out,
largest)) { indices,
largest)) {
// Successed, return. // Successed, return.
return; return;
} else { } else {
...@@ -181,15 +182,16 @@ void TopkKernel(const Context& dev_ctx, ...@@ -181,15 +182,16 @@ void TopkKernel(const Context& dev_ctx,
// The conclusion is drawn from the data through multiple sets of // The conclusion is drawn from the data through multiple sets of
// statistics // statistics
if (input_width >= 128 && k >= input_width * 0.75) { if (input_width >= 128 && k >= input_width * 0.75) {
if (ops::SortTopk<T>( auto* ctx = reinterpret_cast<const paddle::platform::CUDADeviceContext*>(
paddle::platform::CUDADeviceContext(dev_ctx.GetPlace()), &dev_ctx);
&trans_input, if (ops::SortTopk<T>(*ctx,
input_width, &trans_input,
input_height, input_width,
k, input_height,
&trans_out, k,
&trans_ind, &trans_out,
largest)) { &trans_ind,
largest)) {
// last step, tranpose back the indices and output // last step, tranpose back the indices and output
funcs::TransCompute<phi::GPUContext, int64_t>( funcs::TransCompute<phi::GPUContext, int64_t>(
ndims, dev_ctx, trans_ind, indices, trans); ndims, dev_ctx, trans_ind, indices, trans);
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册