diff --git a/core/src/banking_stage/transaction_scheduler/prio_graph_scheduler.rs b/core/src/banking_stage/transaction_scheduler/prio_graph_scheduler.rs
index 6dde96420c1d19..045d2cca1d8dba 100644
--- a/core/src/banking_stage/transaction_scheduler/prio_graph_scheduler.rs
+++ b/core/src/banking_stage/transaction_scheduler/prio_graph_scheduler.rs
@@ -358,12 +358,18 @@ impl PrioGraphScheduler {
     ) {
         let thread_id = self.in_flight_tracker.complete_batch(batch_id);
         for transaction in transactions {
-            let account_locks = transaction.get_account_locks_unchecked();
-            self.account_locks.unlock_accounts(
-                account_locks.writable.into_iter(),
-                account_locks.readonly.into_iter(),
-                thread_id,
-            );
+            let message = transaction.message();
+            let account_keys = message.account_keys();
+            let write_account_locks = account_keys
+                .iter()
+                .enumerate()
+                .filter_map(|(index, key)| message.is_writable(index).then_some(key));
+            let read_account_locks = account_keys
+                .iter()
+                .enumerate()
+                .filter_map(|(index, key)| (!message.is_writable(index)).then_some(key));
+            self.account_locks
+                .unlock_accounts(write_account_locks, read_account_locks, thread_id);
         }
     }