|
|
|
@ -158,11 +158,6 @@ function(op_library TARGET)
|
|
|
|
|
file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL(relu, MKLDNN);\n")
|
|
|
|
|
else()
|
|
|
|
|
file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL(${TARGET}, MKLDNN);\n")
|
|
|
|
|
# HACK: fc only have cpu kernel
|
|
|
|
|
if (${MKLDNN_FILE} STREQUAL "fc_mkldnn_op")
|
|
|
|
|
file(APPEND ${pybind_file} "USE_CPU_ONLY_OP(${TARGET});\n")
|
|
|
|
|
set(pybind_flag 1)
|
|
|
|
|
endif()
|
|
|
|
|
endif()
|
|
|
|
|
endif()
|
|
|
|
|
|
|
|
|
@ -175,6 +170,9 @@ function(op_library TARGET)
|
|
|
|
|
file(APPEND ${pybind_file} "USE_OP(fake_dequantize_max_abs);\n")
|
|
|
|
|
elseif(${TARGET} STREQUAL "tensorrt_engine_op")
|
|
|
|
|
message(STATUS "Pybind skips [tensorrt_engine_op], for this OP is only used in inference")
|
|
|
|
|
elseif(${TARGET} STREQUAL "fc")
|
|
|
|
|
# HACK: fc only have mkldnn and cpu, which would mismatch the cpu only condition
|
|
|
|
|
file(APPEND ${pybind_file} "USE_CPU_ONLY_OP(${TARGET});\n")
|
|
|
|
|
else()
|
|
|
|
|
file(APPEND ${pybind_file} "USE_OP(${TARGET});\n")
|
|
|
|
|
endif()
|
|
|
|
|