diff --git a/custom_components/lanbon_switch/switch.py b/custom_components/lanbon_switch/switch.py
index ea73d01..ba73964 100644
--- a/custom_components/lanbon_switch/switch.py
+++ b/custom_components/lanbon_switch/switch.py
@@ -2,7 +2,7 @@
 from homeassistant.components import mqtt
 from homeassistant.components.switch import SwitchEntity
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN, TOPIC_PREFIX, SWITCH_SUBTOPIC