diff --git a/crates/vm/src/arch/vm.rs b/crates/vm/src/arch/vm.rs
index 8462594d2..ca7c98a64 100644
--- a/crates/vm/src/arch/vm.rs
+++ b/crates/vm/src/arch/vm.rs
@@ -78,7 +78,6 @@ pub struct VmExecutorNextSegmentState<F: PrimeField32> {
     pub memory: MemoryImage<F>,
     pub input: Streams<F>,
     pub pc: u32,
-    pub segment_idx: usize,
 }
 
 pub struct VmExecutorOneSegmentResult<F: PrimeField32, VC: VmConfig<F>> {
@@ -136,8 +135,9 @@ where
         let mut segment_idx = 0;
 
         loop {
+            let _span = info_span!("execute_segment", segment = segment_idx).entered();
             let one_segment_result =
-                self.execute_until_segment(exe.clone(), memory, streams, segment_idx, pc)?;
+                self.execute_until_segment(exe.clone(), memory, streams, pc)?;
             segments.push(one_segment_result.segment);
             if one_segment_result.next_state.is_none() {
                 break;
@@ -145,8 +145,8 @@ where
             let next_state = one_segment_result.next_state.unwrap();
             memory = next_state.memory;
             pc = next_state.pc;
-            segment_idx = next_state.segment_idx;
             streams = next_state.input;
+            segment_idx += 1;
         }
         tracing::debug!("Number of continuation segments: {}", segments.len());
 
@@ -161,7 +161,6 @@ where
         exe: impl Into<VmExe<F>>,
         memory: MemoryImage<F>,
         input: impl Into<Streams<F>>,
-        segment_idx: usize,
         pc: u32,
     ) -> Result<VmExecutorOneSegmentResult<F, VC>, ExecutionError> {
         let exe = exe.into();
@@ -176,7 +175,6 @@ where
         if let Some(overridden_heights) = self.overridden_heights.as_ref() {
             segment.set_override_trace_heights(overridden_heights.clone());
         }
-        let _span = info_span!("execute_segment", segment = segment_idx).entered();
         let state = metrics_span("execute_time_ms", || segment.execute_from_pc(pc))?;
 
         if state.is_terminated {
@@ -205,7 +203,6 @@ where
                 memory: final_memory,
                 input: streams,
                 pc: state.pc,
-                segment_idx: segment_idx + 1,
             }),
         })
     }