diff --git a/train.py b/train.py index 004c8eeda121..4c3bec34835f 100644 --- a/train.py +++ b/train.py @@ -23,7 +23,7 @@ import sys import time from copy import deepcopy -from datetime import datetime +from datetime import datetime, timedelta from pathlib import Path try: @@ -529,7 +529,8 @@ def main(opt, callbacks=Callbacks()): assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command' torch.cuda.set_device(LOCAL_RANK) device = torch.device('cuda', LOCAL_RANK) - dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo') + dist.init_process_group(backend='nccl' if dist.is_nccl_available() else 'gloo', + timeout=timedelta(seconds=10800)) # Train if not opt.evolve: