diff --git a/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py b/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py
index 120723fae5ab..26fbf6164d54 100644
--- a/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py
+++ b/deepspeed/runtime/swap_tensor/partitioned_param_swapper.py
@@ -121,7 +121,7 @@ def _configure_aio(self, ds_config):
                                        dtype=self.dtype,
                                        device=get_accelerator().device_name(),
                                        requires_grad=False)
-            self.aio_read_handle.new_device_locked_tensor(self.buffers)
+            self.aio_read_handle.pin_device_tensor(self.buffers)
         else:
             self.buffers = get_accelerator().pin_memory(torch.empty(int(self.aligned_elements_per_buffer *
                                                                         self.param_buffer_count),