diff --git a/custom_components/badnest/water_heater.py b/custom_components/badnest/water_heater.py
index ca1916a..9d20c06 100644
--- a/custom_components/badnest/water_heater.py
+++ b/custom_components/badnest/water_heater.py
@@ -11,7 +11,6 @@ from homeassistant.const import (
     ATTR_ENTITY_ID,
 )
 from homeassistant.components.water_heater import (
-    WaterHeaterDevice,
     SCAN_INTERVAL,
     STATE_OFF,
     STATE_ON,
@@ -22,6 +21,11 @@ from homeassistant.components.water_heater import (
     ATTR_OPERATION_MODE,
     ATTR_OPERATION_LIST,
 )
+try:
+    from homeassistant.components.water_heater import WaterHeaterEntity
+except ImportError:
+    from homeassistant.components.water_heater import WaterHeaterDevice as WaterHeaterEntity
+
 from .const import (
     DOMAIN,
 )