diff --git a/custom_components/hilo/sensor.py b/custom_components/hilo/sensor.py
index 9522670..a31a934 100755
--- a/custom_components/hilo/sensor.py
+++ b/custom_components/hilo/sensor.py
@@ -1,4 +1,5 @@
 """Support for various Hilo sensors."""
+
 from __future__ import annotations
 
 from datetime import datetime, timedelta, timezone
@@ -24,6 +25,7 @@
     UnitOfTemperature,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.util import Throttle, slugify
@@ -255,7 +257,6 @@ class EnergySensor(IntegrationSensor):
     _attr_native_unit_of_measurement = UnitOfEnergy.KILO_WATT_HOUR
     _attr_state_class = SensorStateClass.TOTAL_INCREASING
     _attr_unit_of_measurement = UnitOfEnergy.KILO_WATT_HOUR
-    _attr_icon = "mdi:lightning-bolt"
 
     def __init__(self, hilo, device):
         self._device = device
@@ -275,6 +276,9 @@ def __init__(self, hilo, device):
         initial_state = 0
         self._attr_native_value = initial_state
         self._attr_last_valid_state = initial_state
+        self._device_info = DeviceInfo(
+            identifiers={(DOMAIN, self._device.identifier)},
+        )
 
         super().__init__(
             integration_method=METHOD_LEFT,
@@ -284,7 +288,9 @@ def __init__(self, hilo, device):
             unique_id=self._attr_unique_id,
             unit_prefix="k",
             unit_time="h",
+            device_info=self._device_info,
         )
+        self._attr_icon = "mdi:lightning-bolt"
         LOG.debug(
             f"Setting up EnergySensor entity: {self._attr_name} with source {self._source}"
         )