diff --git a/score/ommToken/snapshot/snapshot.py b/score/ommToken/snapshot/snapshot.py
index f9db83d9..a76f9ce1 100644
--- a/score/ommToken/snapshot/snapshot.py
+++ b/score/ommToken/snapshot/snapshot.py
@@ -51,7 +51,7 @@ def _createSnapshot(self, _owner: Address, _user_old_staked_balance: int, _user_
         """
         _now = self.now()
         self._snapshot.create_checkpoints(_owner, _now, _user_new_staked_balance)
-        self._snapshot.create_total_checkpoints(_now, _user_new_staked_balance)
+        self._snapshot.create_total_checkpoints(_now, _total_staked_balance)
 
         self.SnapshotCreated(_owner, _user_new_staked_balance, _user_old_staked_balance)
 
diff --git a/score/ommToken/tokens/IRC2.py b/score/ommToken/tokens/IRC2.py
index 9879c1e6..0d108716 100644
--- a/score/ommToken/tokens/IRC2.py
+++ b/score/ommToken/tokens/IRC2.py
@@ -476,3 +476,8 @@ def _mint(self, _to: Address, _amount: int, _data: bytes = None) -> None:
         self._balances[_to] += _amount
 
         self.Transfer(ZERO_SCORE_ADDRESS, _to, _amount, _data)
+
+    @external
+    @only_owner
+    def updateTotalStakedBalanceOfAt(self, _timestamp: int, _staked: int):
+        self._snapshot.create_total_checkpoints(_timestamp, _staked)
\ No newline at end of file