diff --git a/custom_components/zha_toolkit/__init__.py b/custom_components/zha_toolkit/__init__.py
index 20502e9..96121bd 100644
--- a/custom_components/zha_toolkit/__init__.py
+++ b/custom_components/zha_toolkit/__init__.py
@@ -4,7 +4,12 @@
 
 import homeassistant.helpers.config_validation as cv
 import voluptuous as vol
-from homeassistant.components.zha.core.gateway import ZHAGateway
+
+try:
+    from homeassistant.components.zha import Gateway as ZHAGateway
+except ImportError:
+    from homeassistant.components.zha.core.gateway import ZHAGateway
+
 from homeassistant.util import dt as dt_util
 from zigpy import types as t
 from zigpy.exceptions import DeliveryError
diff --git a/custom_components/zha_toolkit/utils.py b/custom_components/zha_toolkit/utils.py
index e3d0307..e897df6 100644
--- a/custom_components/zha_toolkit/utils.py
+++ b/custom_components/zha_toolkit/utils.py
@@ -10,7 +10,12 @@
 from enum import Enum
 
 import zigpy
-from homeassistant.components.zha.core.gateway import ZHAGateway
+
+try:
+    from homeassistant.components.zha import Gateway as ZHAGateway
+except ImportError:
+    from homeassistant.components.zha.core.gateway import ZHAGateway
+
 from homeassistant.util import dt as dt_util
 from pkg_resources import get_distribution, parse_version
 from zigpy import types as t