diff --git a/custom_components/teslemetry/switch.py b/custom_components/teslemetry/switch.py
index 6aa259e..f5d7c3d 100644
--- a/custom_components/teslemetry/switch.py
+++ b/custom_components/teslemetry/switch.py
@@ -19,6 +19,7 @@
 from .const import DOMAIN
 from .entity import (
     TeslemetryVehicleEntity,
+    TeslemetryEnergyInfoEntity
 )
 from .models import TeslemetryVehicleData, TeslemetryEnergyData
 
@@ -169,7 +170,7 @@ async def async_turn_off(self, **kwargs: Any) -> None:
         self.set((self.entity_description.key, False))
 
 
-class TeslemetryEnergyInfoSwitchEntity(TeslemetryVehicleEntity, SwitchEntity):
+class TeslemetryEnergyInfoSwitchEntity(TeslemetryEnergyInfoEntity, SwitchEntity):
     """Base class for Teslemetry Switch."""
 
     _attr_device_class = SwitchDeviceClass.SWITCH