diff --git a/device.mk b/device.mk
index ff449bf..906e9e5 100644
--- a/device.mk
+++ b/device.mk
@@ -258,6 +258,9 @@ PRODUCT_PACKAGES += \
 # Kernel
 PRODUCT_SET_DEBUGFS_RESTRICTIONS := true
 
+# Kernel Headers
+PRODUCT_VENDOR_KERNEL_HEADERS := device/xiaomi/stone-kernel/kernel-headers
+
 # Keymaster
 PRODUCT_PACKAGES += \
     android.hardware.keymaster@4.1.vendor
diff --git a/gpt-utils/Android.bp b/gpt-utils/Android.bp
index b3a5f1c..2705177 100644
--- a/gpt-utils/Android.bp
+++ b/gpt-utils/Android.bp
@@ -35,7 +35,7 @@ cc_library {
     ],
     owner: "qti",
     header_libs: [
-        "generated_kernel_headers",
+        "device_kernel_headers",
     ],
     export_include_dirs: ["."],
 }