diff --git a/checkpoint_engine/ps.py b/checkpoint_engine/ps.py
index 1d8c5cf..7d74eca 100644
--- a/checkpoint_engine/ps.py
+++ b/checkpoint_engine/ps.py
@@ -799,6 +799,7 @@ def _update_per_bucket(
 
         gidx = 0
         ret_code = torch.zeros((), device=self.device_manager.device_type, dtype=torch.int64)
+        buffer_b: torch.Tensor | None = None
         try:
             for i in range(max_len):
                 if i < len(receiver_rank_buckets) and not disable_h2d_buffer: