diff --git a/drivers/thermal/pex9749-thermal.c b/drivers/thermal/pex9749-thermal.c
index a192f534..8c785bbd 100644
--- a/drivers/thermal/pex9749-thermal.c
+++ b/drivers/thermal/pex9749-thermal.c
@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: GPL-2.0
 // Copyright (c) 2022-2023, NVIDIA CORPORATION.  All rights reserved.
 
+#include <nvidia/conftest.h>
+
 #include <linux/module.h>
 #include <linux/i2c.h>
 #include <linux/thermal.h>
@@ -108,7 +110,7 @@ static int pex9749_calc_temp(u16 v)
 	return temp;
 }
 
-#if LINUX_VERSION_CODE >= KERNEL_VERSION(6, 1, 0)
+#if defined(NV_DEVM_THERMAL_OF_ZONE_REGISTER_PRESENT)
 static int pex9749_get_temp(struct thermal_zone_device *tz, int *temp)
 {
 	struct pex9749_priv *priv = tz->devdata;
@@ -156,7 +158,7 @@ out:
 	return ret;
 }
 
-#if LINUX_VERSION_CODE >= KERNEL_VERSION(6, 1, 0)
+#if defined(NV_DEVM_THERMAL_OF_ZONE_REGISTER_PRESENT)
 static struct thermal_zone_device_ops pex9749_ops = {
 #else
 static struct thermal_zone_of_device_ops pex9749_ops = {
@@ -205,7 +207,7 @@ static int pex9749_probe(struct i2c_client *client, const struct i2c_device_id *
 		return -ENOMEM;
 
 	priv->client = client;
-#if LINUX_VERSION_CODE >= KERNEL_VERSION(6, 1, 0)
+#if defined(NV_DEVM_THERMAL_OF_ZONE_REGISTER_PRESENT)
 	priv->tzd = devm_thermal_of_zone_register(dev, PEX9749, priv, &pex9749_ops);
 #else
 	priv->tzd = devm_thermal_zone_of_sensor_register(dev, PEX9749, priv, &pex9749_ops);