diff --git a/custom_components/ohme/base.py b/custom_components/ohme/base.py
index 21523ed..65fe7bd 100644
--- a/custom_components/ohme/base.py
+++ b/custom_components/ohme/base.py
@@ -1,4 +1,4 @@
-from homeassistant.helpers.entity import DeviceInfo, Entity
+from homeassistant.helpers.entity import Entity
 from homeassistant.core import callback
 
 
@@ -17,7 +17,7 @@ def __init__(self, cooordinator, hass, client):
         self._last_updated = None
         self._state = None
 
-        self._attr_device_info = DeviceInfo(**client.get_device_info())
+        self._attr_device_info = client.get_device_info()
 
     async def async_added_to_hass(self) -> None:
         """When entity is added to hass."""