diff --git a/evm/src/generation/prover_input.rs b/evm/src/generation/prover_input.rs
index af3ffe3a..6539feb6 100644
--- a/evm/src/generation/prover_input.rs
+++ b/evm/src/generation/prover_input.rs
@@ -191,11 +191,12 @@ impl<F: Field> GenerationState<F> {
         b_start_loc: usize,
         m_start_loc: usize,
     ) -> (Vec<U256>, Vec<U256>) {
-        let a = &self.memory.contexts[0].segments[Segment::KernelGeneral as usize].content
+        let n = self.memory.contexts.len();
+        let a = &self.memory.contexts[n - 1].segments[Segment::KernelGeneral as usize].content
             [a_start_loc..a_start_loc + len];
-        let b = &self.memory.contexts[0].segments[Segment::KernelGeneral as usize].content
+        let b = &self.memory.contexts[n - 1].segments[Segment::KernelGeneral as usize].content
             [b_start_loc..b_start_loc + len];
-        let m = &self.memory.contexts[0].segments[Segment::KernelGeneral as usize].content
+        let m = &self.memory.contexts[n - 1].segments[Segment::KernelGeneral as usize].content
             [m_start_loc..m_start_loc + len];
 
         let a_biguint = mem_vec_to_biguint(a);