diff --git a/CMakeLists.txt b/CMakeLists.txt
index 74e2b72..0a3895c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -69,15 +69,15 @@ target_include_directories(${PROJECT_NAME}
 target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES})
 target_link_libraries(${PROJECT_NAME} PRIVATE NVHPC::NCCL)
 target_link_libraries(${PROJECT_NAME} PRIVATE NVHPC::MPI_CXX)
-target_link_libraries(${PROJECT_NAME} PRIVATE /opt/yaml-cpp/lib/libyaml-cpp.a)
+target_link_libraries(${PROJECT_NAME} PRIVATE ${YAML_CPP_ROOT}/lib/libyaml-cpp.a)
 
 target_include_directories(${PROJECT_NAME}
-   PRIVATE
-  /opt/yaml-cpp/include
-  ${MPI_CXX_INCLUDE_PATH}
-  ${TORCH_INCLUDE_DIRS}
-  ${NVHPC_CUDA_INCLUDE_DIR}
-  ${NVHPC_NCCL_INCLUDE_DIR}
+    PRIVATE
+    ${YAML_CPP_ROOT}/include
+    ${MPI_CXX_INCLUDE_PATH}
+    ${TORCH_INCLUDE_DIRS}
+    ${NVHPC_CUDA_INCLUDE_DIR}
+    ${NVHPC_NCCL_INCLUDE_DIR}
 )
 target_compile_definitions(${PROJECT_NAME} PRIVATE YAML_CPP_STATIC_DEFINE)
 target_compile_options(${PROJECT_NAME} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${TORCH_CXX_FLAGS}>)
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 7663e04..a685bc6 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -49,6 +49,7 @@ COPY . /torchfort
 RUN cd /torchfort && mkdir build && cd build && \
     cmake -DCMAKE_INSTALL_PREFIX=/opt/torchfort \
     -DNVHPC_CUDA_VERSION=11.8 \
+    -DYAML_CPP_ROOT=/opt/yaml-cpp \
     -DCMAKE_PREFIX_PATH="`python -c 'import torch;print(torch.utils.cmake_prefix_path)'`;/opt/nvidia/hpc_sdk/Linux_x86_64/23.5/cmake" \
     .. && \
     make -j install && \
diff --git a/examples/cpp/cart_pole/CMakeLists.txt b/examples/cpp/cart_pole/CMakeLists.txt
index cade849..a1c0efc 100644
--- a/examples/cpp/cart_pole/CMakeLists.txt
+++ b/examples/cpp/cart_pole/CMakeLists.txt
@@ -22,7 +22,7 @@ target_link_libraries(PyEnvironments PRIVATE environments)
 foreach(tgt ${cart_pole_example_targets})
   target_include_directories(${tgt}
     PRIVATE
-    /opt/yaml-cpp/include
+    ${YAML_CPP_ROOT}/include
     ${MPI_CXX_INCLUDE_PATH}
     ${NVHPC_CUDA_INCLUDE_DIR}
     ${Python_INCLUDE_DIRS}
@@ -32,7 +32,7 @@ foreach(tgt ${cart_pole_example_targets})
   target_link_libraries(${tgt} PRIVATE ${TORCH_LIBRARIES})
   target_link_libraries(${tgt} PRIVATE ${Python_LIBRARIES})
   target_link_libraries(${tgt} PRIVATE NVHPC::MPI_CXX)
-  target_link_libraries(${tgt} PRIVATE /opt/yaml-cpp/lib/libyaml-cpp.a)
+  target_link_libraries(${tgt} PRIVATE ${YAML_CPP_ROOT}/lib/libyaml-cpp.a)
   target_link_libraries(${tgt} PRIVATE environments)
   target_compile_options(${tgt} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${TORCH_CXX_FLAGS}>)
   target_link_options(${tgt} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:${TORCH_CXX_FLAGS}>)