diff --git a/iris-mpc-gpu/src/server/actor.rs b/iris-mpc-gpu/src/server/actor.rs
index 0478a6771..4b90819db 100644
--- a/iris-mpc-gpu/src/server/actor.rs
+++ b/iris-mpc-gpu/src/server/actor.rs
@@ -1212,6 +1212,8 @@ impl ServerActor {
             let now = std::time::Instant::now();
             tracing::info!("Collected enough match distances, starting bucket calculation");
 
+            self.device_manager.await_streams(&batch_streams);
+
             let indices = match_distances_indices
                 .iter()
                 .enumerate()
@@ -1315,6 +1317,8 @@ impl ServerActor {
                 }
             }
 
+            self.device_manager.await_streams(&batch_streams);
+
             tracing::info!("Bucket calculation took {:?}", now.elapsed());
         }