diff --git a/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp b/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
index d21fdbc5f..4982bc910 100644
--- a/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
+++ b/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp
@@ -393,6 +393,11 @@ hsa_status_t XdnaDriver::ConfigHwCtxCU(
   amdxdna_hwctx_param_config_cu *xdna_config_cu_param =
       reinterpret_cast<amdxdna_hwctx_param_config_cu *>(
           malloc(config_cu_param_size));
+  if (xdna_config_cu_param == nullptr) {
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+  MAKE_SCOPE_GUARD([xdna_config_cu_param] { free(xdna_config_cu_param); });
+
   xdna_config_cu_param->num_cus = config_cu_param.num_cus;
 
   for (int i = 0; i < xdna_config_cu_param->num_cus; ++i) {