use determined device for device handle

pytorch · May 17, 2023 · dd032d5 · dd032d5
1 parent 8ac84cd
commit dd032d5
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/torch/distributed/fsdp/_init_utils.py b/torch/distributed/fsdp/_init_utils.py
@@ -304,9 +304,7 @@ def _init_device_handle(
             "cuda", torch.cuda.current_device()
         )
 
-    state._device_handle = _FSDPDeviceHandle.from_device(
-        torch.device("cuda", torch.cuda.current_device())
-    )
+    state._device_handle = _FSDPDeviceHandle.from_device(determined_device)
     return state