diff --git a/deploy/stage/smpcv2-0-stage/values-upgrade-server-left.yaml b/deploy/stage/smpcv2-0-stage/values-upgrade-server-left.yaml
index 90c4fbd21..81b44cf68 100644
--- a/deploy/stage/smpcv2-0-stage/values-upgrade-server-left.yaml
+++ b/deploy/stage/smpcv2-0-stage/values-upgrade-server-left.yaml
@@ -9,8 +9,6 @@ args:
   - "left"
   - "--environment"
   - "$(ENVIRONMENT)"
-  - "--batch-size"
-  - "50"
 
 initContainer:
   enabled: true
diff --git a/deploy/stage/smpcv2-0-stage/values-upgrade-server-right.yaml b/deploy/stage/smpcv2-0-stage/values-upgrade-server-right.yaml
index aae283072..45f690e25 100644
--- a/deploy/stage/smpcv2-0-stage/values-upgrade-server-right.yaml
+++ b/deploy/stage/smpcv2-0-stage/values-upgrade-server-right.yaml
@@ -9,8 +9,6 @@ args:
   - "right"
   - "--environment"
   - "$(ENVIRONMENT)"
-  - "--batch-size"
-  - "50"
 
 initContainer:
   enabled: true
diff --git a/deploy/stage/smpcv2-1-stage/values-upgrade-server-left.yaml b/deploy/stage/smpcv2-1-stage/values-upgrade-server-left.yaml
index 26248cb14..4cb1c1b49 100644
--- a/deploy/stage/smpcv2-1-stage/values-upgrade-server-left.yaml
+++ b/deploy/stage/smpcv2-1-stage/values-upgrade-server-left.yaml
@@ -9,8 +9,6 @@ args:
   - "left"
   - "--environment"
   - "$(ENVIRONMENT)"
-  - "--batch-size"
-  - "50"
 
 initContainer:
   enabled: true
diff --git a/deploy/stage/smpcv2-1-stage/values-upgrade-server-right.yaml b/deploy/stage/smpcv2-1-stage/values-upgrade-server-right.yaml
index 90603205e..9ec8a96b2 100644
--- a/deploy/stage/smpcv2-1-stage/values-upgrade-server-right.yaml
+++ b/deploy/stage/smpcv2-1-stage/values-upgrade-server-right.yaml
@@ -9,8 +9,6 @@ args:
   - "right"
   - "--environment"
   - "$(ENVIRONMENT)"
-  - "--batch-size"
-  - "50"
 
 initContainer:
   enabled: true
diff --git a/deploy/stage/smpcv2-2-stage/values-upgrade-server-left.yaml b/deploy/stage/smpcv2-2-stage/values-upgrade-server-left.yaml
index a744738ba..2f181eea7 100644
--- a/deploy/stage/smpcv2-2-stage/values-upgrade-server-left.yaml
+++ b/deploy/stage/smpcv2-2-stage/values-upgrade-server-left.yaml
@@ -9,8 +9,6 @@ args:
   - "left"
   - "--environment"
   - "$(ENVIRONMENT)"
-  - "--batch-size"
-  - "50"
 
 initContainer:
   enabled: true
diff --git a/deploy/stage/smpcv2-2-stage/values-upgrade-server-right.yaml b/deploy/stage/smpcv2-2-stage/values-upgrade-server-right.yaml
index 3f8f226b1..2d1f948bb 100644
--- a/deploy/stage/smpcv2-2-stage/values-upgrade-server-right.yaml
+++ b/deploy/stage/smpcv2-2-stage/values-upgrade-server-right.yaml
@@ -9,8 +9,6 @@ args:
   - "right"
   - "--environment"
   - "$(ENVIRONMENT)"
-  - "--batch-size"
-  - "50"
 
 initContainer:
   enabled: true
diff --git a/iris-mpc-upgrade/src/bin/tcp_ssl_upgrade_client.rs b/iris-mpc-upgrade/src/bin/tcp_ssl_upgrade_client.rs
index cb40c4c4c..537a7546e 100644
--- a/iris-mpc-upgrade/src/bin/tcp_ssl_upgrade_client.rs
+++ b/iris-mpc-upgrade/src/bin/tcp_ssl_upgrade_client.rs
@@ -87,10 +87,15 @@ async fn main() -> eyre::Result<()> {
     server1.write_u64(start).await?;
     server2.write_u64(start).await?;
     server3.write_u64(start).await?;
+
     server1.write_u64(end).await?;
     server2.write_u64(end).await?;
     server3.write_u64(end).await?;
 
+    server1.write_u8(args.batch_size).await?;
+    server2.write_u8(args.batch_size).await?;
+    server3.write_u8(args.batch_size).await?;
+
     server1.flush().await?;
     server2.flush().await?;
     server3.flush().await?;
diff --git a/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs b/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
index 88bbe63ee..92311b868 100644
--- a/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
+++ b/iris-mpc-upgrade/src/bin/tcp_upgrade_server.rs
@@ -107,22 +107,26 @@ async fn main() -> eyre::Result<()> {
             end2
         );
     }
+
+    let batch_size1 = client_stream1.read_u8().await?;
+    let batch_size2 = client_stream2.read_u8().await?;
+
+    if batch_size1 != batch_size2 {
+        bail!(
+            "Invalid batch size: client1: {}, client2: {}",
+            batch_size1,
+            batch_size2,
+        );
+    }
+
     let num_elements = end1.checked_sub(start1).unwrap();
-    let num_batches = num_elements / u64::from(args.batch_size);
-    tracing::info!(
-        "Batch size: {}, num batches: {}",
-        args.batch_size,
-        num_batches
-    );
+    let num_batches = num_elements / u64::from(batch_size1);
+    tracing::info!("Batch size: {}, num batches: {}", batch_size1, num_batches);
 
     let mut batch = Vec::new();
 
     for batch_num in 0..num_batches + 1 {
-        tracing::info!(
-            "Processing batch {} of size: {}",
-            batch_num,
-            args.batch_size
-        );
+        tracing::info!("Processing batch {} of size: {}", batch_num, batch_size1);
         let start_time = Instant::now();
         let batch_size_1_message = client_stream1.read_u8().await?;
         let batch_size_2_message = client_stream2.read_u8().await?;
diff --git a/iris-mpc-upgrade/src/config.rs b/iris-mpc-upgrade/src/config.rs
index 15df6e325..453f7ba81 100644
--- a/iris-mpc-upgrade/src/config.rs
+++ b/iris-mpc-upgrade/src/config.rs
@@ -49,9 +49,6 @@ pub struct UpgradeServerConfig {
     #[clap(long)]
     pub party_id: PartyID,
 
-    #[clap(long)]
-    pub batch_size: u8,
-
     #[clap(long)]
     pub eye: Eye,
 
@@ -59,7 +56,7 @@ pub struct UpgradeServerConfig {
     pub environment: String,
 }
 
-impl std::fmt::Debug for UpgradeServerConfig {
+impl fmt::Debug for UpgradeServerConfig {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.debug_struct("UpgradeServerConfig")
             .field("bind_addr", &self.bind_addr)