Deprecations fix `TORCH_NCCL_BLOCKING_WAIT` (#9448)

10 months ago · 7df821e6ea
parent aa756069d4
commit 7df821e6ea
1 changed files with 1 additions and 1 deletions
--- a/ultralytics/engine/trainer.py
+++ b/ultralytics/engine/trainer.py
@ -210,7 +210,7 @@ class BaseTrainer:
        torch.cuda.set_device(RANK)
        self.device = torch.device("cuda", RANK)
        # LOGGER.info(f'DDP info: RANK {RANK}, WORLD_SIZE {world_size}, DEVICE {self.device}')
-        os.environ["NCCL_BLOCKING_WAIT"] = "1"  # set to enforce timeout
+        os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1"  # set to enforce timeout
        dist.init_process_group(
            "nccl" if dist.is_nccl_available() else "gloo",
            timeout=timedelta(seconds=10800),  # 3 hours