diff --git a/core/src/main/java/kafka/autobalancer/model/ClusterModel.java b/core/src/main/java/kafka/autobalancer/model/ClusterModel.java
index e7e1de9467..68cb234dfa 100644
--- a/core/src/main/java/kafka/autobalancer/model/ClusterModel.java
+++ b/core/src/main/java/kafka/autobalancer/model/ClusterModel.java
@@ -321,8 +321,7 @@ public void reassignPartition(Uuid topicId, int partitionId, int brokerId) {
             }
             if (oldBrokerId != -1) {
                 TopicPartition tp = new TopicPartition(topicName, partitionId);
-                TopicPartitionReplicaUpdater replicaUpdater = brokerReplicaMap.get(oldBrokerId).get(tp);
-                brokerReplicaMap.get(brokerId).put(tp, replicaUpdater);
+                brokerReplicaMap.get(brokerId).put(tp, createReplicaUpdater(tp));
                 brokerReplicaMap.get(oldBrokerId).remove(tp);
             }
             topicPartitionReplicaMap.get(topicName).put(partitionId, brokerId);