diff --git a/.github/workflows/temp-branch-build-and-push.yaml b/.github/workflows/temp-branch-build-and-push.yaml
index 07e966a44..694a0632e 100644
--- a/.github/workflows/temp-branch-build-and-push.yaml
+++ b/.github/workflows/temp-branch-build-and-push.yaml
@@ -3,7 +3,7 @@ name: Branch - Build and push docker image
 on:
   push:
     branches:
-      - "test-benchmark"
+      - "fix/test-benchmark-chunked-page-lock"
 
 concurrency:
   group: '${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}'
diff --git a/iris-mpc-gpu/src/helpers/mod.rs b/iris-mpc-gpu/src/helpers/mod.rs
index e757e924b..9a7a505d9 100644
--- a/iris-mpc-gpu/src/helpers/mod.rs
+++ b/iris-mpc-gpu/src/helpers/mod.rs
@@ -183,17 +183,18 @@ pub fn register_host_memory(
         chunk_offset + chunk_length
     );
     let size = chunk_length / device_manager.device_count();
+    let offset = chunk_offset / device_manager.device_count();
     for (device_index, device) in device_manager.devices().iter().enumerate() {
         device.bind_to_thread().unwrap();
         unsafe {
             let _ = cudarc::driver::sys::lib().cuMemHostRegister_v2(
-                (db.limb_0[device_index] + (chunk_offset * code_length) as u64) as *mut _,
+                (db.limb_0[device_index] + (offset * code_length) as u64) as *mut _,
                 size * code_length,
                 CU_MEMHOSTALLOC_PORTABLE,
             );
 
             let _ = cudarc::driver::sys::lib().cuMemHostRegister_v2(
-                (db.limb_1[device_index] + (chunk_offset * code_length) as u64) as *mut _,
+                (db.limb_1[device_index] + (offset * code_length) as u64) as *mut _,
                 size * code_length,
                 CU_MEMHOSTALLOC_PORTABLE,
             );