diff --git a/locked-asset/factory/src/lib.rs b/locked-asset/factory/src/lib.rs
index a758faa50..8ccc375ae 100644
--- a/locked-asset/factory/src/lib.rs
+++ b/locked-asset/factory/src/lib.rs
@@ -147,7 +147,7 @@ pub trait LockedAssetFactory:
     #[endpoint(unlockAssets)]
     fn unlock_assets(&self) {
         self.require_not_paused();
-        
+
         let (token_id, token_nonce, amount) = self.call_value().single_esdt().into_tuple();
         let locked_token_id = self.locked_asset_token_id().get();
         require!(token_id == locked_token_id, "Bad payment token");
@@ -199,10 +199,14 @@ pub trait LockedAssetFactory:
         self.send()
             .esdt_local_burn(&locked_token_id, token_nonce, &amount);
 
-        let amounts_per_epoch = attributes.get_unlock_amounts_per_epoch(&amount);
-        let unlockable_amounts_per_epoch =
-            amounts_per_epoch.get_unlockable_entries(month_start_epoch);
-        self.update_energy_after_unlock(caller.clone(), unlockable_amounts_per_epoch);
+        let initial_amounts_per_epoch = attributes.get_unlock_amounts_per_epoch(&amount);
+        let final_amounts_per_epoch =
+            output_locked_asset_attributes.get_unlock_amounts_per_epoch(&locked_remaining);
+        self.update_energy_after_unlock(
+            caller.clone(),
+            initial_amounts_per_epoch,
+            final_amounts_per_epoch,
+        );
 
         self.emit_unlock_assets_event(
             &caller,
diff --git a/locked-asset/factory/src/migration.rs b/locked-asset/factory/src/migration.rs
index 9e0732c4e..2f6c9643f 100644
--- a/locked-asset/factory/src/migration.rs
+++ b/locked-asset/factory/src/migration.rs
@@ -13,7 +13,8 @@ mod energy_factory_proxy {
         fn update_energy_after_old_token_unlock(
             &self,
             original_caller: ManagedAddress,
-            epoch_amount_pairs: UnlockEpochAmountPairs<Self::Api>,
+            initial_epoch_amount_pairs: UnlockEpochAmountPairs<Self::Api>,
+            final_epoch_amount_pairs: UnlockEpochAmountPairs<Self::Api>,
         );
     }
 }
@@ -35,12 +36,17 @@ pub trait LockedTokenMigrationModule:
     fn update_energy_after_unlock(
         &self,
         caller: ManagedAddress,
-        epoch_amount_pairs: UnlockEpochAmountPairs<Self::Api>,
+        initial_epoch_amount_pairs: UnlockEpochAmountPairs<Self::Api>,
+        final_epoch_amount_pairs: UnlockEpochAmountPairs<Self::Api>,
     ) {
         let new_factory_address = self.new_factory_address().get();
         let _: IgnoreValue = self
             .new_factory_proxy_builder(new_factory_address)
-            .update_energy_after_old_token_unlock(caller, epoch_amount_pairs)
+            .update_energy_after_old_token_unlock(
+                caller,
+                initial_epoch_amount_pairs,
+                final_epoch_amount_pairs,
+            )
             .execute_on_dest_context();
     }