diff --git a/custom_components/hilo/__init__.py b/custom_components/hilo/__init__.py
index d0ebe48..1d68185 100755
--- a/custom_components/hilo/__init__.py
+++ b/custom_components/hilo/__init__.py
@@ -1,4 +1,5 @@
 """Support for Hilo automation systems."""
+
 from __future__ import annotations
 
 import asyncio
@@ -240,7 +241,8 @@ def __init__(self, hass: HomeAssistant, entry: ConfigEntry, api: API) -> None:
             CONF_APPRECIATION_PHASE, DEFAULT_APPRECIATION_PHASE
         )
         self.pre_cold = entry.options.get(
-            CONF_PRE_COLD_PHASE, DEFAULT_PRE_COLD_PHASE  # this is new
+            CONF_PRE_COLD_PHASE,
+            DEFAULT_PRE_COLD_PHASE,
         )
         self.challenge_lock = entry.options.get(
             CONF_CHALLENGE_LOCK, DEFAULT_CHALLENGE_LOCK
@@ -774,7 +776,7 @@ def __init__(
         hilo: Hilo,
         name: Union[str, None] = None,
         *,
-        device: HiloDevice | None = None,
+        device: HiloDevice,
     ) -> None:
         """Initialize."""
         assert hilo.coordinator
diff --git a/custom_components/hilo/switch.py b/custom_components/hilo/switch.py
index 73020f3..6456393 100644
--- a/custom_components/hilo/switch.py
+++ b/custom_components/hilo/switch.py
@@ -4,6 +4,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
+from pyhilo.device.switch import Switch
 
 from . import Hilo, HiloEntity
 from .const import DOMAIN, LOG, SWITCH_CLASSES
@@ -23,7 +24,7 @@ async def async_setup_entry(
 
 
 class HiloSwitch(HiloEntity, SwitchEntity):
-    def __init__(self, hilo: Hilo, device):
+    def __init__(self, hilo: Hilo, device: Switch):
         super().__init__(hilo, device=device, name=device.name)
         old_unique_id = f"{slugify(device.name)}-switch"
         self._attr_unique_id = f"{slugify(device.identifier)}-switch"