diff --git a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
index c8c4d8458..3b3e8da6b 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
@@ -60,9 +60,6 @@ public class InFlightRequestTracker {
   public void removeBatch(int batchId, String hostAndPushPort) {
     Set<Integer> batchIdSet = inflightBatchesPerAddress.get(hostAndPushPort);
     batchIdSet.remove(batchId);
-    if (batchIdSet.size() == 0) {
-      inflightBatchesPerAddress.remove(hostAndPushPort);
-    }
   }
 
   public void onSuccess(String hostAndPushPort) {