Skip to content

Commit e7f0696

Browse files
committed
Top K with omp
1 parent 3d6934e commit e7f0696

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

paddle/fluid/operators/top_k_op.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,9 @@ class TopkKernel : public framework::OpKernel<T> {
5555
// NOTE: eigen shape doesn't affect paddle tensor.
5656
eg_input.reshape(flat2dims);
5757

58+
#ifdef PADDLE_WITH_MKLDNN
59+
#pragma omp parallel for
60+
#endif
5861
for (size_t i = 0; i < row; i++) {
5962
std::vector<std::pair<T, size_t>> vec;
6063
for (size_t j = 0; j < col; j++) {

0 commit comments

Comments
 (0)