diff --git a/paddle/pten/api/lib/CMakeLists.txt b/paddle/pten/api/lib/CMakeLists.txt
index 189548880694d9c413cf306f9995a49fc70d1b0e..33f20a4df04135a91e07f78550bc543062d3a292 100644
--- a/paddle/pten/api/lib/CMakeLists.txt
+++ b/paddle/pten/api/lib/CMakeLists.txt
@@ -24,6 +24,7 @@ set(api_source_file_tmp ${api_source_file}.tmp)
 
 add_custom_command(
   OUTPUT ${api_header_file} ${api_source_file}
+  COMMAND ${PYTHON_EXECUTABLE} -m pip install pyyaml
   COMMAND ${PYTHON_EXECUTABLE} ${api_gen_file} 
                  --api_yaml_path ${api_yaml_file}
                  --api_header_path ${api_header_file_tmp}