diff --git a/scripts/build_local.sh b/scripts/build_local.sh
index 23e1394dfff..ea684de04d6 100755
--- a/scripts/build_local.sh
+++ b/scripts/build_local.sh
@@ -31,6 +31,12 @@ if [ -n "${USE_ANACONDA}" ]; then
   export CONDA_INSTALL_LOCALLY=1
   "${ROOT_DIR}/scripts/build_anaconda.sh" "$@"
 else
+  # Make sure that pyyaml is installed for the codegen of building Aten to work
+  if [[ -n "$(python -c 'import yaml' 2>&1)" ]]; then
+    echo "Installing pyyaml with pip at $(which pip)"
+    pip install --user pyyaml
+  fi
+
   # Build protobuf compiler from third_party if configured to do so
   if [ -n "${USE_HOST_PROTOC:-}" ]; then
     echo "USE_HOST_PROTOC is set; building protoc before building Caffe2..."
diff --git a/scripts/build_windows.bat b/scripts/build_windows.bat
index 185d43ac0a8..c4b8a38fafa 100644
--- a/scripts/build_windows.bat
+++ b/scripts/build_windows.bat
@@ -37,6 +37,9 @@ if NOT DEFINED CMAKE_GENERATOR (
   )
 )
 
+:: Install pyyaml for Aten codegen
+pip install pyyaml
+
 echo CAFFE2_ROOT=%CAFFE2_ROOT%
 echo CMAKE_GENERATOR=%CMAKE_GENERATOR%
 echo CMAKE_BUILD_TYPE=%CMAKE_BUILD_TYPE%