Use device-agnostic accelerator API.

Signed-off-by: Jagadish Krishnamoorthy <[email protected]>
deepspeedai · Oct 15, 2024 · 7f91988 · 7f91988
1 parent 3d679a5
commit 7f91988
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/training/cifar/cifar10_deepspeed.py b/training/cifar/cifar10_deepspeed.py
@@ -281,8 +281,7 @@ def main(args):
     # Initialize DeepSpeed distributed backend.
     deepspeed.init_distributed()
     _local_rank = int(os.environ.get("LOCAL_RANK"))
-    torch_device = torch.device(f"cuda:{_local_rank}")
-    torch.cuda.set_device(torch_device)
+    get_accelerator().set_device(_local_rank)
 
     ########################################################################
     # Step1. Data Preparation.