From b11846789b09c30563e7e841a3e1577c12bb07ed Mon Sep 17 00:00:00 2001 From: Nadav Elyahu Date: Wed, 27 Mar 2024 09:30:11 +0200 Subject: [PATCH] Fix wrong device access in get_lst_from_rank0 --- deepspeed/runtime/zero/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/zero/utils.py b/deepspeed/runtime/zero/utils.py index 3993689294c7..e26004c605c3 100755 --- a/deepspeed/runtime/zero/utils.py +++ b/deepspeed/runtime/zero/utils.py @@ -65,7 +65,7 @@ def get_lst_from_rank0(lst: List[int]) -> None: lst_tensor = torch.tensor( lst if dist.get_rank() == 0 else [-1] * len(lst), dtype=int, - device=get_accelerator().device(os.environ["LOCAL_RANK"]), + device=torch.device(get_accelerator().device_name(os.environ["LOCAL_RANK"])), requires_grad=False, ) dist.broadcast(lst_tensor, src=0, async_op=False)