diff --git a/paddle/fluid/operators/top_k_op.h b/paddle/fluid/operators/top_k_op.h index d44eeae8e6ff9ac87ab093d04e3f5427743f0c08..b31fee22a2d21ac9f2c79546cee3a266fb0ed1d9 100644 --- a/paddle/fluid/operators/top_k_op.h +++ b/paddle/fluid/operators/top_k_op.h @@ -55,6 +55,9 @@ class TopkKernel : public framework::OpKernel { // NOTE: eigen shape doesn't affect paddle tensor. eg_input.reshape(flat2dims); +#ifdef PADDLE_WITH_MKLDNN +#pragma omp parallel for +#endif for (size_t i = 0; i < row; i++) { std::vector> vec; for (size_t j = 0; j < col; j++) {