diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index f1394b62b8255..91ff6e4052a81 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -1971,6 +1971,9 @@ if(BUILD_PYTHON) target_include_directories(caffe2_pybind11_state_gpu PRIVATE ${Caffe2_CPU_INCLUDE}) target_link_libraries(caffe2_pybind11_state_gpu PRIVATE torch_library python::python pybind::pybind11) + if(USE_MKLDNN) + target_link_libraries(caffe2_pybind11_state_gpu PRIVATE caffe2::mkldnn) + endif() if(WIN32) target_link_libraries(caffe2_pybind11_state_gpu PRIVATE onnx_proto) endif(WIN32)