diff --git a/custom_components/rte_ecowatt/__init__.py b/custom_components/rte_ecowatt/__init__.py
index 9395578..0bfd327 100644
--- a/custom_components/rte_ecowatt/__init__.py
+++ b/custom_components/rte_ecowatt/__init__.py
@@ -562,8 +562,7 @@ async def async_migrate_entry(hass, config_entry: ConfigEntry):
         )
         new = {**config_entry.data}
         new["enedis_load_shedding"] = [False]
-        config_entry.version = 2
-        hass.config_entries.async_update_entry(config_entry, data=new)
+        hass.config_entries.async_update_entry(config_entry, data=new, version=2)
         _LOGGER.info(f"Migration to version {config_entry.version} successful")
     if config_entry.version < 3:
         _LOGGER.warn(
@@ -571,7 +570,6 @@ async def async_migrate_entry(hass, config_entry: ConfigEntry):
         )
         new = {**config_entry.data}
         del new["enedis_load_shedding"]
-        config_entry.version = 3
-        hass.config_entries.async_update_entry(config_entry, data=new)
+        hass.config_entries.async_update_entry(config_entry, data=new, version=3)
         _LOGGER.info(f"Migration to version {config_entry.version} successful")
     return True