diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt
index f5ef2fb51f..a0f1131122 100644
--- a/cmake/CMakeLists.txt
+++ b/cmake/CMakeLists.txt
@@ -1632,11 +1632,14 @@ set(onnxruntime_DELAYLOAD_FLAGS "")
 include_directories(
   ${ONNXRUNTIME_INCLUDE_DIR}
   ${REPO_ROOT}/include/onnxruntime/core/session
-  if (onnxruntime_ENABLE_TRAINING_ON_DEVICE)
-    ${REPO_ROOT}/orttraining/orttraining/training_api/include/
-  endif()
 )
 
+if (onnxruntime_ENABLE_TRAINING_ON_DEVICE)
+  include_directories(
+    ${REPO_ROOT}/orttraining/orttraining/training_api/include/
+  )
+endif()
+
 if (onnxruntime_USE_OPENVINO)
 
   add_definitions(-DUSE_OPENVINO=1)
diff --git a/orttraining/orttraining/training_api/module.cc b/orttraining/orttraining/training_api/module.cc
index 85201afdc5..634629c8b9 100644
--- a/orttraining/orttraining/training_api/module.cc
+++ b/orttraining/orttraining/training_api/module.cc
@@ -109,7 +109,7 @@ Module::Module(const std::string& train_model_path_or_bytes,
     ORT_ENFORCE(params_iter != named_parameters_.end());
 
     // Retrieve the target device for "param_name"
-    std::vector<SessionState::NodeInfo> node_info_vec;
+    InlinedVector<SessionState::NodeInfo> node_info_vec;
     ORT_THROW_IF_ERROR(train_sess_state.GetInputNodeInfo(param_name, node_info_vec));
     const auto& node_info = node_info_vec.front();
     const auto target_device = *node_info.device;