diff --git a/core/blockchain.go b/core/blockchain.go
index fd337edd0..864ead777 100644
--- a/core/blockchain.go
+++ b/core/blockchain.go
@@ -1159,6 +1159,12 @@ func (bc *BlockChain) WriteBlockWithState(
 	if len(header.ShardState()) > 0 {
 		// Write shard state for the new epoch
 		epoch := new(big.Int).Add(header.Epoch(), common.Big1)
+		shardState, err := block.Header().GetShardState()
+		if err == nil && shardState.Epoch != nil && bc.chainConfig.IsStaking(shardState.Epoch) {
+			// After staking, the epoch will be decided by the epoch in the shard state.
+			epoch = new(big.Int).Set(shardState.Epoch)
+		}
+
 		newShardState, err := bc.WriteShardStateBytes(batch, epoch, header.ShardState())
 		if err != nil {
 			header.Logger(utils.Logger()).Warn().Err(err).Msg("cannot store shard state")
diff --git a/internal/utils/utils.go b/internal/utils/utils.go
index c4c73e7c8..50878283b 100644
--- a/internal/utils/utils.go
+++ b/internal/utils/utils.go
@@ -95,8 +95,9 @@ func GetAddressFromBlsPubKeyBytes(pubKeyBytes []byte) common.Address {
 	if err == nil {
 		addrBytes := pubKey.GetAddress()
 		addr.SetBytes(addrBytes[:])
+	} else {
+		Logger().Err(err).Msg("Failed to get address of bls key")
 	}
-	Logger().Err(err).Msg("Failed to get address of bls key")
 	return addr
 }