Update

zeshengzong · zeshengzong · commit 01b1d4238016 · 2025-03-14T10:57:12.000+08:00
diff --git a/torch/utils/data/_utils/pin_memory.py b/torch/utils/data/_utils/pin_memory.py
@@ -21,16 +21,7 @@ def _pin_memory_loop(in_queue, out_queue, device_id, done_event, device):
     torch.set_num_threads(1)
 
     torch.multiprocessing._set_thread_name("pt_data_pin")
-
-    if device == "cuda":
-        torch.cuda.set_device(device_id)
-    elif device == "xpu":
-        torch.xpu.set_device(device_id)  # type: ignore[attr-defined]
-    elif device == torch._C._get_privateuse1_backend_name():
-        custom_device_mod = getattr(torch, torch._C._get_privateuse1_backend_name())
-        custom_device_mod.set_device(device_id)
-    elif device is None:
-        torch.accelerator.set_device_index(device_id)
+    torch.accelerator.set_device_index(device_id)
 
     def do_one_step():
         try:
