diff --git a/xsense/async_xsense.py b/xsense/async_xsense.py
index c2cc493..7520d96 100644
--- a/xsense/async_xsense.py
+++ b/xsense/async_xsense.py
@@ -238,7 +238,7 @@ async def get_state(self, station: Station):
             raise APIFailure(f'Unable to retrieve station data: {self._lastres.status}/{text}')
 
     async def set_state(self, entity: Entity, shadow: str, topic: str, definition: Dict):
-        station = entity.station
+        station = entity if isinstance(entity, Station) else entity.station
         t = datetime.now()
         timestamp = t.strftime('%Y%m%d%H%M%S')