[Hardware] Replace torch.cuda.device_count/current_device/set_device API (#36145)

Signed-off-by: Kunshang Ji <jikunshang95@gmail.com> Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-03-12 22:57:47 +08:00
parent 2e693f48e7
commit 53ec16a705
89 changed files with 254 additions and 219 deletions
--- a/benchmarks/kernels/benchmark_device_communicators.py
+++ b/benchmarks/kernels/benchmark_device_communicators.py
@@ -495,7 +495,7 @@ def main():

    # Set device
    device = torch.device(f"cuda:{rank}")
-    torch.cuda.set_device(device)
+    torch.accelerator.set_device_index(device)

    # Get CPU process group
    cpu_group = dist.new_group(backend="gloo")