diff --git a/vm/src/vm/runners/builtin_runner/output.rs b/vm/src/vm/runners/builtin_runner/output.rs
index a0b6559014..7b6c54f920 100644
--- a/vm/src/vm/runners/builtin_runner/output.rs
+++ b/vm/src/vm/runners/builtin_runner/output.rs
@@ -170,6 +170,21 @@ impl OutputBuiltinRunner {
 
         Ok(())
     }
+
+    pub fn get_public_memory(&self) -> Result<Vec<(usize, usize)>, RunnerError> {
+        let size = self
+            .stop_ptr
+            .ok_or(RunnerError::NoStopPointer(Box::new(OUTPUT_BUILTIN_NAME)))?;
+
+        let mut public_memory: Vec<(usize, usize)> = (0..size).map(|i| (i, 0)).collect();
+        for (page_id, page) in self.pages.iter() {
+            for index in 0..page.size {
+                public_memory[page.start + index].1 = page_id.clone();
+            }
+        }
+
+        Ok(public_memory)
+    }
 }
 
 impl Default for OutputBuiltinRunner {
@@ -551,4 +566,22 @@ mod tests {
             matches!(result, Err(RunnerError::PageNotOnSegment(relocatable, base)) if relocatable == page_start && base == builtin.base())
         )
     }
+
+    #[test]
+    fn get_public_memory() {
+        let mut builtin = OutputBuiltinRunner::new(true);
+        let page_start = Relocatable {
+            segment_index: builtin.base() as isize,
+            offset: 2,
+        };
+        builtin
+            .add_page(1, page_start.clone(), 2)
+            .expect("Failed to add page");
+
+        // Mock the effects of `final_stack()`
+        builtin.stop_ptr = Some(4);
+
+        let public_memory = builtin.get_public_memory().unwrap();
+        assert_eq!(public_memory, vec![(0, 0), (1, 0), (2, 1), (3, 1)]);
+    }
 }
diff --git a/vm/src/vm/runners/cairo_runner.rs b/vm/src/vm/runners/cairo_runner.rs
index fb7b44b07c..ee3e050478 100644
--- a/vm/src/vm/runners/cairo_runner.rs
+++ b/vm/src/vm/runners/cairo_runner.rs
@@ -56,7 +56,7 @@ use num_traits::{ToPrimitive, Zero};
 use serde::{Deserialize, Serialize};
 
 use super::{
-    builtin_runner::{KeccakBuiltinRunner, PoseidonBuiltinRunner, OUTPUT_BUILTIN_NAME},
+    builtin_runner::{KeccakBuiltinRunner, PoseidonBuiltinRunner},
     cairo_pie::{self, CairoPie, CairoPieMetadata, CairoPieVersion},
 };
 
@@ -1091,8 +1091,8 @@ impl CairoRunner {
             let (_, size) = builtin_runner
                 .get_used_cells_and_allocated_size(vm)
                 .map_err(RunnerError::FinalizeSegements)?;
-            if builtin_runner.name() == OUTPUT_BUILTIN_NAME {
-                let public_memory = (0..size).map(|i| (i, 0)).collect();
+            if let BuiltinRunner::Output(output_builtin) = builtin_runner {
+                let public_memory = output_builtin.get_public_memory()?;
                 vm.segments
                     .finalize(Some(size), builtin_runner.base(), Some(&public_memory))
             } else {